def _get_ar_order(cls, rhos: TensorVariable, ar_order: Optional[int], constant: bool) -> int:
"""Compute ar_order given inputs
If ar_order is not specified we do constant folding on the shape of rhos
to retrieve it. For example, this will detect that
Normal(size=(5, 3)).shape[-1] == 3, which is not known by Aesara before.
Raises
------
ValueError
If inferred ar_order cannot be inferred from rhos or if it is less than 1
"""
if ar_order is None:
shape_fg = FunctionGraph(
outputs=[rhos.shape[-1]],
features=[ShapeFeature()],
clone=True,
)
(folded_shape,) = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs
folded_shape = getattr(folded_shape, "data", None)
if folded_shape is None:
raise ValueError(
"Could not infer ar_order from last dimension of rho. Pass it "
"explictily or make sure rho have a static shape"
)
ar_order = int(folded_shape) - int(constant)
if ar_order < 1:
raise ValueError(
"Inferred ar_order is smaller than 1. Increase the last dimension "
"of rho or remove constant_term"
)
return ar_order
def test_mvnormal_ShapeFeature():
M_aet = iscalar("M")
M_aet.tag.test_value = 2
d_rv = multivariate_normal(aet.ones((M_aet, )), aet.eye(M_aet), size=2)
fg = FunctionGraph(
[i for i in graph_inputs([d_rv]) if not isinstance(i, Constant)],
[d_rv],
clone=False,
features=[ShapeFeature()],
)
s1, s2 = fg.shape_feature.shape_of[d_rv]
assert get_test_value(s1) == 2
assert M_aet in graph_inputs([s2])
# Test broadcasted shapes
mean = tensor(config.floatX, [True, False])
mean.tag.test_value = np.array([[0, 1, 2]], dtype=config.floatX)
test_covar = np.diag(np.array([1, 10, 100], dtype=config.floatX))
test_covar = np.stack([test_covar, test_covar * 10.0])
cov = aet.as_tensor(test_covar).type()
cov.tag.test_value = test_covar
d_rv = multivariate_normal(mean, cov, size=[2, 3])
fg = FunctionGraph(
[i for i in graph_inputs([d_rv]) if not isinstance(i, Constant)],
[d_rv],
clone=False,
features=[ShapeFeature()],
)
s1, s2, s3, s4 = fg.shape_feature.shape_of[d_rv]
assert s1.get_test_value() == 2
assert s2.get_test_value() == 3
assert s3.get_test_value() == 2
assert s4.get_test_value() == 3
def infer_shape(outs, inputs, input_shapes):
"""
Compute the shape of the outputs given the shape of the inputs of an Aesara
graph.
We do it this way to avoid compiling the inner function just to get
the shape. Changes to ShapeFeature could require changes in this function.
"""
# We use a ShapeFeature because it has all the necessary logic
# inside. We don't use the full ShapeFeature interface, but we
# let it initialize itself with an empty fgraph, otherwise we will
# need to do it manually
for inp, inp_shp in zip(inputs, input_shapes):
if inp_shp is not None and len(inp_shp) != inp.type.ndim:
assert len(inp_shp) == inp.type.ndim
shape_feature = ShapeFeature()
shape_feature.on_attach(FunctionGraph([], []))
# Initialize shape_of with the input shapes
for inp, inp_shp in zip(inputs, input_shapes):
shape_feature.set_shape(inp, inp_shp)
def local_traverse(out):
"""
Go back in the graph, from out, adding computable shapes to shape_of.
"""
if out in shape_feature.shape_of:
# Its shape is already known
return
elif out.owner is None:
# This is an input of the graph
shape_feature.init_r(out)
else:
# Recurse over inputs
for inp in out.owner.inputs:
if inp not in shape_feature.shape_of:
local_traverse(inp)
# shape_feature.on_import does not actually use an fgraph
# It will call infer_shape and set_shape appropriately
dummy_fgraph = None
shape_feature.on_import(dummy_fgraph, out.owner, reason="dummy")
ret = []
for o in outs:
local_traverse(o)
ret.append(shape_feature.shape_of[o])
return ret
def test_normal_ShapeFeature():
M_aet = iscalar("M")
M_aet.tag.test_value = 3
sd_aet = scalar("sd")
sd_aet.tag.test_value = np.array(1.0, dtype=config.floatX)
d_rv = normal(aet.ones((M_aet, )), sd_aet, size=(2, M_aet))
d_rv.tag.test_value
fg = FunctionGraph(
[i for i in graph_inputs([d_rv]) if not isinstance(i, Constant)],
[d_rv],
clone=False,
features=[ShapeFeature()],
)
s1, s2 = fg.shape_feature.shape_of[d_rv]
assert get_test_value(s1) == get_test_value(d_rv).shape[0]
assert get_test_value(s2) == get_test_value(d_rv).shape[1]
def test_dirichlet_ShapeFeature():
"""Make sure `RandomVariable.infer_shape` works with `ShapeFeature`."""
M_at = iscalar("M")
M_at.tag.test_value = 2
N_at = iscalar("N")
N_at.tag.test_value = 3
d_rv = dirichlet(at.ones((M_at, N_at)), name="Gamma")
fg = FunctionGraph(
outputs=[d_rv],
clone=False,
features=[ShapeFeature()],
)
s1, s2 = fg.shape_feature.shape_of[d_rv]
assert M_at in graph_inputs([s1])
assert N_at in graph_inputs([s2])
def test_dirichlet_ShapeFeature():
"""Make sure `RandomVariable.infer_shape` works with `ShapeFeature`."""
M_tt = iscalar("M")
M_tt.tag.test_value = 2
N_tt = iscalar("N")
N_tt.tag.test_value = 3
d_rv = dirichlet(aet.ones((M_tt, N_tt)), name="Gamma")
fg = FunctionGraph(
[i for i in graph_inputs([d_rv]) if not isinstance(i, Constant)],
[d_rv],
clone=False,
features=[ShapeFeature()],
)
s1, s2 = fg.shape_feature.shape_of[d_rv]
assert M_tt in graph_inputs([s1])
assert N_tt in graph_inputs([s2])
def test_nonstandard_shapes():
a = tensor3(config.floatX)
a.tag.test_value = np.random.random((2, 3, 4)).astype(config.floatX)
b = tensor3(config.floatX)
b.tag.test_value = np.random.random((2, 3, 4)).astype(config.floatX)
tl = make_list([a, b])
tl_shape = shape(tl)
assert np.array_equal(tl_shape.get_test_value(), (2, 2, 3, 4))
# There's no `FunctionGraph`, so it should return a `Subtensor`
tl_shape_i = shape_i(tl, 0)
assert isinstance(tl_shape_i.owner.op, Subtensor)
assert tl_shape_i.get_test_value() == 2
tl_fg = FunctionGraph([a, b], [tl], features=[ShapeFeature()])
tl_shape_i = shape_i(tl, 0, fgraph=tl_fg)
assert not isinstance(tl_shape_i.owner.op, Subtensor)
assert tl_shape_i.get_test_value() == 2
none_shape = shape(NoneConst)
assert np.array_equal(none_shape.get_test_value(), [])
