def tt_broadcast_arrays(*args: TensorVariable):
"""Broadcast any number of arrays against each other.
Parameters
----------
`*args` : array_likes
The arrays to broadcast.
"""
bcast_shape = broadcast_shape(*args)
return tuple(at_broadcast_to(a, bcast_shape) for a in args)
def test_broadcast_shape_symbolic(s1_vals, s2_vals, exp_res):
s1s = aet.lscalars(len(s1_vals))
eval_point = {}
for s, s_val in zip(s1s, s1_vals):
eval_point[s] = s_val
s.tag.test_value = s_val
s2s = aet.lscalars(len(s2_vals))
for s, s_val in zip(s2s, s2_vals):
eval_point[s] = s_val
s.tag.test_value = s_val
res = broadcast_shape(s1s, s2s, arrays_are_shapes=True)
res = aet.as_tensor(res)
assert tuple(res.eval(eval_point)) == exp_res
def test_broadcast_shape():
def shape_tuple(x, use_bcast=True):
if use_bcast:
return tuple(s if not bcast else 1
for s, bcast in zip(tuple(x.shape), x.broadcastable))
else:
return tuple(s for s in tuple(x.shape))
x = np.array([[1], [2], [3]])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# Now, we try again using shapes as the inputs
#
# This case also confirms that a broadcast dimension will
# broadcast against a non-broadcast dimension when they're
# both symbolic (i.e. we couldn't obtain constant values).
b_aet = broadcast_shape(
shape_tuple(x_aet, use_bcast=False),
shape_tuple(y_aet, use_bcast=False),
arrays_are_shapes=True,
)
assert any(
isinstance(node.op, Assert)
for node in applys_between([x_aet, y_aet], b_aet))
assert np.array_equal([z.eval() for z in b_aet], b.shape)
b_aet = broadcast_shape(shape_tuple(x_aet),
shape_tuple(y_aet),
arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# These are all constants, so there shouldn't be any asserts in the
# resulting graph.
assert not any(
isinstance(node.op, Assert)
for node in applys_between([x_aet, y_aet], b_aet))
x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
b_aet = broadcast_shape(shape_tuple(x_aet),
shape_tuple(y_aet),
arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
x = np.empty((1, 2, 3))
y = np.array(1)
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert b_aet[0].value == 1
assert np.array_equal([z.eval() for z in b_aet], b.shape)
assert not any(
isinstance(node.op, Assert)
for node in applys_between([x_aet, y_aet], b_aet))
b_aet = broadcast_shape(shape_tuple(x_aet),
shape_tuple(y_aet),
arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
x = np.empty((2, 1, 3))
y = np.empty((2, 1, 1))
b = np.broadcast(x, y)
x_aet = aet.as_tensor_variable(x)
y_aet = aet.as_tensor_variable(y)
b_aet = broadcast_shape(x_aet, y_aet)
assert b_aet[1].value == 1
assert np.array_equal([z.eval() for z in b_aet], b.shape)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
b_aet = broadcast_shape(shape_tuple(x_aet),
shape_tuple(y_aet),
arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_aet], b.shape)
x1_shp_aet = iscalar("x1")
x2_shp_aet = iscalar("x2")
y1_shp_aet = iscalar("y1")
x_shapes = (1, x1_shp_aet, x2_shp_aet)
x_aet = aet.ones(x_shapes)
y_shapes = (y1_shp_aet, 1, x2_shp_aet)
y_aet = aet.ones(y_shapes)
b_aet = broadcast_shape(x_aet, y_aet)
# TODO: This will work when/if we use a more sophisticated `is_same_graph`
# implementation.
# assert not any(
# isinstance(node.op, Assert)
# for node in graph_ops([x_aet, y_aet], b_aet)
# )
res = aet.as_tensor(b_aet).eval({
x1_shp_aet: 10,
x2_shp_aet: 4,
y1_shp_aet: 2,
})
assert np.array_equal(res, (2, 10, 4))
y_shapes = (y1_shp_aet, 1, y1_shp_aet)
y_aet = aet.ones(y_shapes)
b_aet = broadcast_shape(x_aet, y_aet)
assert isinstance(b_aet[-1].owner.op, Assert)
def test_broadcast_shape_basic():
def shape_tuple(x, use_bcast=True):
if use_bcast:
return tuple(
s if not bcast else 1
for s, bcast in zip(tuple(x.shape), x.broadcastable)
)
else:
return tuple(s for s in tuple(x.shape))
x = np.array([[1], [2], [3]])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_at = at.as_tensor_variable(x)
y_at = at.as_tensor_variable(y)
b_at = broadcast_shape(x_at, y_at)
assert np.array_equal([z.eval() for z in b_at], b.shape)
# Now, we try again using shapes as the inputs
#
# This case also confirms that a broadcast dimension will
# broadcast against a non-broadcast dimension when they're
# both symbolic (i.e. we couldn't obtain constant values).
b_at = broadcast_shape(
shape_tuple(x_at, use_bcast=False),
shape_tuple(y_at, use_bcast=False),
arrays_are_shapes=True,
)
assert any(
isinstance(node.op, Assert) for node in applys_between([x_at, y_at], b_at)
)
assert np.array_equal([z.eval() for z in b_at], b.shape)
b_at = broadcast_shape(shape_tuple(x_at), shape_tuple(y_at), arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_at], b.shape)
x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
b = np.broadcast(x, y)
x_at = at.as_tensor_variable(x)
y_at = at.as_tensor_variable(y)
b_at = broadcast_shape(x_at, y_at)
assert np.array_equal([z.eval() for z in b_at], b.shape)
b_at = broadcast_shape(shape_tuple(x_at), shape_tuple(y_at), arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_at], b.shape)
x = np.empty((1, 2, 3))
y = np.array(1)
b = np.broadcast(x, y)
x_at = at.as_tensor_variable(x)
y_at = at.as_tensor_variable(y)
b_at = broadcast_shape(x_at, y_at)
assert b_at[0].value == 1
assert np.array_equal([z.eval() for z in b_at], b.shape)
b_at = broadcast_shape(shape_tuple(x_at), shape_tuple(y_at), arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_at], b.shape)
x = np.empty((2, 1, 3))
y = np.empty((2, 1, 1))
b = np.broadcast(x, y)
x_at = at.as_tensor_variable(x)
y_at = at.as_tensor_variable(y)
b_at = broadcast_shape(x_at, y_at)
assert b_at[1].value == 1
assert np.array_equal([z.eval() for z in b_at], b.shape)
b_at = broadcast_shape(shape_tuple(x_at), shape_tuple(y_at), arrays_are_shapes=True)
assert np.array_equal([z.eval() for z in b_at], b.shape)
x1_shp_at = iscalar("x1")
x2_shp_at = iscalar("x2")
y1_shp_at = iscalar("y1")
x_shapes = (1, x1_shp_at, x2_shp_at)
x_at = at.ones(x_shapes)
y_shapes = (y1_shp_at, 1, x2_shp_at)
y_at = at.ones(y_shapes)
b_at = broadcast_shape(x_at, y_at)
res = at.as_tensor(b_at).eval(
{
x1_shp_at: 10,
x2_shp_at: 4,
y1_shp_at: 2,
}
)
assert np.array_equal(res, (2, 10, 4))
y_shapes = (y1_shp_at, 1, y1_shp_at)
y_at = at.ones(y_shapes)
b_at = broadcast_shape(x_at, y_at)
assert isinstance(b_at[-1].owner.op, Assert)
