I have three tensors, A, B and C
in tensorflow, A
and B
are both of shape (m, n, r)
, C
is a binary tensor of sha
import tensorflow as tf
def broadcast(tensor, shape):
"""Broadcasts ``x`` to have shape ``shape``.
|
Uses ``tf.Assert`` statements to ensure that the broadcast is
valid.
First calculates the number of missing dimensions in
``tf.shape(x)`` and left-pads the shape of ``x`` with that many
ones. Then identifies the dimensions of ``x`` that require
tiling and tiles those dimensions appropriately.
Args:
x (tf.Tensor): The tensor to broadcast.
shape (Union[tf.TensorShape, tf.Tensor, Sequence[int]]):
The shape to broadcast to.
Returns:
tf.Tensor: ``x``, reshaped and tiled to have shape ``shape``.
"""
with tf.name_scope('broadcast') as scope:
shape_x = tf.shape(x)
rank_x = tf.shape(shape0)[0]
shape_t = tf.convert_to_tensor(shape, preferred_dtype=tf.int32)
rank_t = tf.shape(shape1)[0]
with tf.control_dependencies([tf.Assert(
rank_t >= rank_x,
['len(shape) must be >= tf.rank(x)', shape_x, shape_t],
summarize=255
)]):
missing_dims = tf.ones(tf.stack([rank_t - rank_x], 0), tf.int32)
shape_x_ = tf.concat([missing_dims, shape_x], 0)
should_tile = tf.equal(shape_x_, 1)
with tf.control_dependencies([tf.Assert(
tf.reduce_all(tf.logical_or(tf.equal(shape_x_, shape_t), should_tile),
['cannot broadcast shapes', shape_x, shape_t],
summarize=255
)]):
multiples = tf.where(should_tile, shape_t, tf.ones_like(shape_t))
out = tf.tile(tf.reshape(x, shape_x_), multiples, name=scope)
try:
out.set_shape(shape)
except:
pass
return out
A = tf.random_normal([20, 100, 10])
B = tf.random_normal([20, 100, 10])
C = tf.random_normal([20, 100, 1])
C = broadcast(C, A.shape)
D = tf.select(C, A, B)