How can I implement this layer similar to Conv2D using tensorflow?

Question

I want to make a neural network layer similar to Conv2D using tensorflow . Below is what I want to implement. A layer uses a kernel just like convolution layer but the output is larger than the input.

The layer image that I want to implement

However, it seems there is no way I can implement that using only tensorflow operations. I managed to implement the below code by converting tensorflow tensors to numpy arrays but I still have no idea how to merge 4D output array into 2D array.

input = [[a, b],
         [c, d]]
kernel = [[1, -1],
          [2, 1]]
output = [[input[0][0] * kernel, input[0][1] * kernel],
          [input[1][0] * kernel, input[1][1] * kernel]]

#since "input[0][0] * kernel" is 2D, "output" becomes 4D array.

Is there any way I can implement this using only tensorflow ? If not, what method should I use instead?

Answer 1

class MyLayer(tf.keras.layers.Layer):
    def __init__(self, kernel):
        super(MyLayer, self).__init__()
        self.k = tf.constant(kernel)

    def build(self, input_shape):
        self.i = input_shape

    def call(self, input):
        x = tf.reshape(input, [-1])
        return tf.map_fn(lambda s:  tf.scalar_mul(s, self.k), x)

mylayer = MyLayer([[1.0, -1.0], [-1.0, 1.0]])
x = tf.constant([[1.0, 2.0, 3.0], [3.0, 4.0, 5.0]])

with tf.Session() as sess:
    print (sess.run(r))

Output:

[[[ 1. -1.]
  [-1.  1.]]
 [[ 2. -2.]
  [-2.  2.]]
 [[ 3. -3.]
  [-3.  3.]]
 [[ 3. -3.]
  [-3.  3.]]
 [[ 4. -4.]
  [-4.  4.]]
 [[ 5. -5.]
  [-5.  5.]]]