Explicit broadcasting of variable batch-size tensor

Question

I'm trying to implement a custom Keras Layer in Tensorflow 2.0RC and need to concatenate a [None, Q] shaped tensor onto a [None, H, W, D] shaped tensor to produce a [None, H, W, D + Q] shaped tensor. It is assumed that the two input tensors have the same batch size even though it is not known beforehand. Also, none of H, W, D, and Q are known at write-time but are evaluated in the layer's build method when the layer is first called. The issue that I'm experiencing is when broadcasting the [None, Q] shaped tensor up to a [None, H, W, Q] shaped tensor in order to concatenate.

Here is an example of trying to create a Keras Model using the Functional API that performs variable-batch broadcasting from shape [None, 3] to shape [None, 5, 5, 3]:

import tensorflow as tf
import tensorflow.keras.layers as kl
import numpy as np

x = tf.keras.Input([3])  # Shape [None, 3]
y = kl.Reshape([1, 1, 3])(x)  # Need to add empty dims before broadcasting
y = tf.broadcast_to(y, [-1, 5, 5, 3])  # Broadcast to shape [None, 5, 5, 3]

model = tf.keras.Model(inputs=x, outputs=y)

print(model(np.random.random(size=(8, 3))).shape)

Tensorflow produces the error:

InvalidArgumentError:  Dimension -1 must be >= 0

And then when I change -1 to None it gives me:

TypeError: Failed to convert object of type <class 'list'> to Tensor. Contents: [None, 5, 5, 3]. Consider casting elements to a supported type.

How can I perform the specified broadcasting?

Answer 1

You need to use the dynamic shape of y to determine the batch size. The dynamic shape of a tensor y is given by tf.shape(y) and is a tensor op representing the shape of y evaluated at runtime. The modified example demonstrates this by selecting between the old shape, [None, 1, 1, 3], and the new shape using tf.where.

import tensorflow as tf
import tensorflow.keras.layers as kl
import numpy as np

x = tf.keras.Input([3])  # Shape [None, 3]
y = kl.Reshape([1, 1, 3])(x)  # Need to add empty dims before broadcasting
# Retain the batch and depth dimensions, but broadcast along H and W
broadcast_shape = tf.where([True, False, False, True],
                           tf.shape(y), [0, 5, 5, 0])
y = tf.broadcast_to(y, broadcast_shape)  # Broadcast to shape [None, 5, 5, 3]

model = tf.keras.Model(inputs=x, outputs=y)

print(model(np.random.random(size=(8, 3))).shape)
# prints: "(8, 5, 5, 3)"

References:

"TensorFlow: Shapes and dynamic dimensions"