def lift_rv_shapes(node):
"""Lift `RandomVariable`'s shape-related parameters.
In other words, this will broadcast the distribution parameters and
extra dimensions added by the `size` parameter.
For example, ``normal([0.0, 1.0], 5.0, size=(3, 2))`` becomes
``normal([[0., 1.], [0., 1.], [0., 1.]], [[5., 5.], [5., 5.], [5., 5.]])``.
"""
if not isinstance(node.op, RandomVariable):
return False
rng, size, dtype, *dist_params = node.inputs
dist_params = broadcast_params(dist_params, node.op.ndims_params)
dist_params = [
broadcast_to(p, (tuple(size) +
tuple(p.shape)) if node.op.ndim_supp > 0 else size)
for p in dist_params
]
return node.op.make_node(rng, None, dtype, *dist_params)
def test_perform(self):
a = tt.scalar()
a.tag.test_value = 5
s_1 = tt.iscalar("s_1")
s_1.tag.test_value = 4
shape = (s_1, 1)
bcast_res = broadcast_to(a, shape)
assert bcast_res.broadcastable == (False, True)
bcast_np = np.broadcast_to(5, (4, 1))
bcast_tt = bcast_res.get_test_value()
assert np.array_equal(bcast_tt, bcast_np)
assert np.shares_memory(bcast_tt, a.get_test_value())
class TestBroadcastTo(utt.InferShapeTester):
rng = np.random.RandomState(43)
def setup_method(self):
super().setup_method()
self.op_class = BroadcastTo
self.op = broadcast_to
@config.change_flags(compute_test_value="raise")
def test_perform(self):
a = tt.scalar()
a.tag.test_value = 5
s_1 = tt.iscalar("s_1")
s_1.tag.test_value = 4
shape = (s_1, 1)
bcast_res = broadcast_to(a, shape)
assert bcast_res.broadcastable == (False, True)
bcast_np = np.broadcast_to(5, (4, 1))
bcast_tt = bcast_res.get_test_value()
assert np.array_equal(bcast_tt, bcast_np)
assert np.shares_memory(bcast_tt, a.get_test_value())
@pytest.mark.parametrize(
"fn,input_dims",
[
[lambda x: broadcast_to(x, (1,)), (1,)],
[lambda x: broadcast_to(x, (6, 2, 5, 3)), (1,)],
[lambda x: broadcast_to(x, (6, 2, 5, 3)), (5, 1)],
[lambda x: broadcast_to(x, (6, 2, 1, 3)), (2, 1, 3)],
],
)
def test_gradient(self, fn, input_dims):
utt.verify_grad(
fn,
[np.random.rand(*input_dims).astype(config.floatX)],
n_tests=1,
rng=self.rng,
)
def test_infer_shape(self):
a = tt.tensor(config.floatX, [False, True, False])
shape = list(a.shape)
out = self.op(a, shape)
self._compile_and_check(
[a] + shape,
[out],
[np.random.rand(2, 1, 3).astype(config.floatX), 2, 1, 3],
self.op_class,
)
a = tt.tensor(config.floatX, [False, True, False])
shape = [tt.iscalar() for i in range(4)]
self._compile_and_check(
[a] + shape,
[self.op(a, shape)],
[np.random.rand(2, 1, 3).astype(config.floatX), 6, 2, 5, 3],
self.op_class,
)