def local_rv_size_lift(fgraph, node):
"""Lift the ``size`` parameter in a ``RandomVariable``.
In other words, this will broadcast the distribution parameters by adding
the extra dimensions implied by the ``size`` parameter, and remove the
``size`` parameter in the process.
For example, ``normal(0, 1, size=(1, 2))`` becomes
``normal([[0, 0]], [[1, 1]], size=())``.
"""
if not isinstance(node.op, RandomVariable):
return
rng, size, dtype, *dist_params = node.inputs
dist_params = broadcast_params(dist_params, node.op.ndims_params)
if get_vector_length(size) > 0:
dist_params = [
broadcast_to(p,
(tuple(size) +
tuple(p.shape)) if node.op.ndim_supp > 0 else size)
for p in dist_params
]
else:
return
new_node = node.op.make_node(rng, None, dtype, *dist_params)
if config.compute_test_value != "off":
compute_test_value(new_node)
return new_node.outputs
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_avoid_useless_subtensors(self):
x = scalar()
y = broadcast_to(x, (1, 2))
# There shouldn't be any unnecessary `Subtensor` operations
# (e.g. from `at.as_tensor((1, 2))[0]`)
assert y.owner.inputs[1].owner is None
assert y.owner.inputs[2].owner is None
def test_perform(self):
a = scalar()
a.tag.test_value = 5
s_1 = 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_aet = bcast_res.get_test_value()
assert np.array_equal(bcast_aet, bcast_np)
assert np.shares_memory(bcast_aet, a.get_test_value())
def test_inplace(self):
"""Make sure that in-place optimizations are *not* performed on the output of a ``BroadcastTo``."""
a = aet.zeros((5, ))
d = aet.vector("d")
c = aet.set_subtensor(a[np.r_[0, 1, 3]], d)
b = broadcast_to(c, (5, ))
q = b[np.r_[0, 1, 3]]
e = aet.set_subtensor(q, np.r_[0, 0, 0])
opts = Query(include=["inplace"])
py_mode = Mode("py", opts)
e_fn = function([d], e, mode=py_mode)
advincsub_node = e_fn.maker.fgraph.outputs[0].owner
assert isinstance(advincsub_node.op, AdvancedIncSubtensor1)
assert isinstance(advincsub_node.inputs[0].owner.op, BroadcastTo)
assert advincsub_node.op.inplace is False
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 = scalar()
a.tag.test_value = 5
s_1 = 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_aet = bcast_res.get_test_value()
assert np.array_equal(bcast_aet, bcast_np)
assert np.shares_memory(bcast_aet, 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 = 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 = tensor(config.floatX, [False, True, False])
shape = [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,
)
def test_inplace(self):
"""Make sure that in-place optimizations are *not* performed on the output of a ``BroadcastTo``."""
a = aet.zeros((5, ))
d = aet.vector("d")
c = aet.set_subtensor(a[np.r_[0, 1, 3]], d)
b = broadcast_to(c, (5, ))
q = b[np.r_[0, 1, 3]]
e = aet.set_subtensor(q, np.r_[0, 0, 0])
opts = Query(include=["inplace"])
py_mode = Mode("py", opts)
e_fn = function([d], e, mode=py_mode)
advincsub_node = e_fn.maker.fgraph.outputs[0].owner
assert isinstance(advincsub_node.op, AdvancedIncSubtensor1)
assert isinstance(advincsub_node.inputs[0].owner.op, BroadcastTo)
assert advincsub_node.op.inplace is False
def test_avoid_useless_scalars(self):
x = scalar()
y = broadcast_to(x, ())
assert y is x
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 = scalar()
a.tag.test_value = 5
s_1 = 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_aet = bcast_res.get_test_value()
assert np.array_equal(bcast_aet, bcast_np)
assert np.shares_memory(bcast_aet, 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 = 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 = tensor(config.floatX, [False, True, False])
shape = [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,
)
def test_BroadcastTo(x, shape):
out = at_extra_ops.broadcast_to(x, shape)
fgraph = FunctionGraph(outputs=[out])
compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
