Copy tensor using K.tile()

I have tensor (None, 196) and after reshaping, it becomes (None, 14, 14).
And now, I want to copy it to channel axis, so that the shape should be (None, 14, 14, 512). Lastly, I want to copy to timestep axis, so it becomes (None, 10, 14, 14, 512). I accomplish those steps using this snippet code:

def replicate(tensor, input_target):
  batch_size = K.shape(tensor)[0]
  nf, h, w, c = input_target
  x = K.reshape(tensor, [batch_size, 1, h, w, 1])
        
  # Replicate to channel dimension
  x = K.tile(x, [batch_size, 1, 1, 1, c])
        
  # Replicate to timesteps dimension
  x = K.tile(x, [batch_size, nf, 1, 1, 1])

  return x

x = ...
x = Lambda(replicate, arguments={'input_target':input_shape})(x)
another_x = Input(shape=input_shape) # shape (10, 14, 14, 512)
        
x = layers.multiply([x, another_x])
x = ...

I plot the model and the output shape is just like I want it to be. But, the problem arises in model training. I set the batch size to 2. This the the error message:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [8,10,14,14,512] vs. [2,10,14,14,512]
[[{{node multiply_1/mul}} = Mul[T=DT_FLOAT, _class=["loc:@training/Adam/gradients/multiply_1/mul_grad/Sum"], _device="/job:localhost/replica:0/task:0/device:GPU:0"](Lambda_2/Tile_1, _arg_another_x_0_0/_189)]]
[[{{node metrics/top_k_categorical_accuracy/Mean_1/_265}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_6346_metrics/top_k_categorical_accuracy/Mean_1", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Looks like, K.tile() increases the batch size from 2 to 8. When I set the batch size to 10, it becomes 1000.

So, my question is how to achieve the result as I want? Is it good way to use tile()? Or, should I use repeat_elements()? Thanks!

I am using Tensorflow 1.12.0 and Keras 2.2.4.

As a rule of thumb, try to avoid bringing batch size to the transformations happening in the Lambda layer.

When you use tile operation, you only set only the dimension that needs to change (for example you had batch_size value in your tile operation which is wrong). Also I am using tf.tile instead of K.tile (TF 1.12 doesn't have tile in the Keras backend it seems).

def replicate(tensor, input_target):
  
  _, nf, h, w, c = input_target
  x = K.reshape(tensor, [-1, 1, h, w, 1])

  # Replicate to channel dimension
  # You can combine below lines to tf.tile(x, [1, nf, 1, 1, c]) as well
  x = tf.tile(x, [1, 1, 1, 1, c])
  # Replicate to timesteps dimension
  x = tf.tile(x, [1, nf, 1, 1, 1])
  

  return x

Simple example

input_shape= [None, 10, 14, 14, 512]
x = Input(shape=(196,))
x = Lambda(replicate, arguments={'input_target':input_shape})(x)
print(x.shape)

Which gives

>>> (?, 10, 14, 14, 512)