How to do math on tensorflow placeholders with one unknown dimension

Question

I need to do math on a tensorflow placeholder, or the data that is passed into it, with shape (None, 128, 128, 3). I need to add a filter with shape (None, 5, 5, 3) to this placeholder, at location [:, i:i+5, j:j+5, 3]. How would I do this?

Before I used the data with length None, I would use

outs = tf.tensor_scatter_nd_add(outs, [[[i + k, j + l] for k in range(5)] for l in
                                                                            range(5)], self.b[h])

nested in two for loops, where outs is the placeholder, self.b[h] was the filter, and i and j indexes from a loop.

Answer 1

Creating a Keras layer:

class BatchAdd(keras.layers.Layer):
    def __init__(self, i, j):
        super(BatchAdd, self).__init__()
        self.i = i
        self.j = j
        self.add_filter = add_filter
    def call(self, outs, b):
        output = tf.vectorized_map(add_filter,
                                   elems=[outs, b, tf.repeat(self.i, tf.shape(b)[0]), tf.repeat(self.j, tf.shape(b)[0])])
        return output

Creating the model

outs = keras.Input(shape=(128, 128, 3))
b = keras.Input(shape=(5, 5, 3))
output = BatchAdd(i,j)(outs, b)
model = keras.Model(inputs=(outs, b), outputs=output)

Check for any batch size:

batch_size = 3
model((tf.random.normal((batch_size, 128, 128, 3)),tf.random.normal((batch_size, 5, 5, 3))))

#output shape: 
   shape=(3, 128, 128, 3)

---

You can use tf.vectorized_map

i = 5 
j = 9

def add_filter(x):
    return tf.tensor_scatter_nd_add(x[0], [[[x[2] + k, x[3] + l] for k in range(5)] for l in
                                                                            range(5)], x[1])
output = tf.vectorized_map(
    add_filter,
    elems=[outs, b, tf.repeat(i, tf.shape(b)[0]), tf.repeat(j, tf.shape(b)[0])])