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(), [])