Использовать tf.shape(target) для получения тензора с формой другого тензора, размер которого может меняться от шага к шагу.
Tensorflow, кстати, умеет автоматически приводить формы. Если сложить тензор А с формой (K, L, M) и тензор B с формой (1, L, M), то TF автоматом скопирует B по первой оси K раз, а при сложении тензоров с формами (K, 1) и (1, L) тоже скопирует оба и результат будет формы (K, L):
https://www.tensorflow.org/xla/broadcasting