def test_RandomVariable_bcast():
rv = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True)
mu = tensor(config.floatX, [True, False, False])
mu.tag.test_value = np.zeros((1, 2, 3)).astype(config.floatX)
sd = tensor(config.floatX, [False, False])
sd.tag.test_value = np.ones((2, 3)).astype(config.floatX)
s1 = iscalar()
s1.tag.test_value = 1
s2 = iscalar()
s2.tag.test_value = 2
s3 = iscalar()
s3.tag.test_value = 3
s3 = Assert("testing")(s3, eq(s1, 1))
res = rv(mu, sd, size=(s1, s2, s3))
assert res.broadcastable == (False, ) * 3
size = aet.as_tensor((1, 2, 3), dtype=np.int32).astype(np.int64)
res = rv(mu, sd, size=size)
assert res.broadcastable == (True, False, False)
res = rv(0, 1, size=aet.as_tensor(1, dtype=np.int64))
assert res.broadcastable == (True, )
def get_steps(
steps: Optional[Union[int, np.ndarray, TensorVariable]],
*,
shape: Optional[Shape] = None,
dims: Optional[Dims] = None,
observed: Optional[Any] = None,
step_shape_offset: int = 0,
):
"""Extract number of steps from shape / dims / observed information
Parameters
----------
steps:
User specified steps for timeseries distribution
shape:
User specified shape for timeseries distribution
dims:
User specified dims for timeseries distribution
observed:
User specified observed data from timeseries distribution
step_shape_offset:
Difference between last shape dimension and number of steps in timeseries
distribution, defaults to 0
Returns
-------
steps
Steps, if specified directly by user, or inferred from the last dimension of
shape / dims / observed. When two sources of step information are provided,
a symbolic Assert is added to ensure they are consistent.
"""
inferred_steps = None
if shape is not None:
shape = to_tuple(shape)
if shape[-1] is not ...:
inferred_steps = shape[-1] - step_shape_offset
if inferred_steps is None and dims is not None:
dims = convert_dims(dims)
if dims[-1] is not ...:
model = modelcontext(None)
inferred_steps = model.dim_lengths[dims[-1]] - step_shape_offset
if inferred_steps is None and observed is not None:
observed = convert_observed_data(observed)
inferred_steps = observed.shape[-1] - step_shape_offset
if inferred_steps is None:
inferred_steps = steps
# If there are two sources of information for the steps, assert they are consistent
elif steps is not None:
inferred_steps = Assert(msg="Steps do not match last shape dimension")(
inferred_steps, at.eq(inferred_steps, steps)
)
return inferred_steps
def test_RandomVariable_bcast_specify_shape():
rv = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True)
s1 = aet.as_tensor(1, dtype=np.int64)
s2 = iscalar()
s2.tag.test_value = 2
s3 = iscalar()
s3.tag.test_value = 3
s3 = Assert("testing")(s3, eq(s1, 1))
size = specify_shape(aet.as_tensor([s1, s3, s2, s2, s1]), (5, ))
mu = tensor(config.floatX, [False, False, True])
mu.tag.test_value = np.random.normal(size=(2, 2, 1)).astype(config.floatX)
std = tensor(config.floatX, [False, True, True])
std.tag.test_value = np.ones((2, 1, 1)).astype(config.floatX)
res = rv(mu, std, size=size)
assert res.broadcastable == (True, False, False, False, True)