I need to input a pooling module to the LSTM decoder, and I'm constructing this using a custom layer with the encoder LSTM states and Keras Input layer as inputs. In this custom layer, I need to scatter the updates to the indices:
updates: <tf.Tensor --- shape=(None, 225, 5, 32) dtype=float32>
indices: <tf.Tensor --- shape=(None, 225) dtype=int32>
with tf.scatter_nd to create a tensor with shape=(None, 960, 5, 32) , something like this:
tf.scatter_nd(tf.expand_dims(indices, 2), updates, shape=[None, 960, 5, 32])
but the problem is that doing this, rises error due to NoneType in shape and I don't want to declare the batch_size in it because it is a Keras layer and only is certain in the learning process. in this state the working version of code is this:
tf.scatter_nd(tf.expand_dims(indices, 2), updates, shape=[960, 5, 32])
>>> <tf.Tensor 'ScatterNd_4:0' shape=(960, 5, 32) dtype=float32>
that has ignored the batch_size in the output.
Is there any alternative way to construct the needed output tensor instead of tf.scatter_nd or a way to make this work properly?