diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index bb45bca96d..c67f70e2e0 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -10,10 +10,10 @@ ### New Features - The `CAR` distribution has been added to allow for use of conditional autoregressions which often are used in spatial and network models. - The dimensionality of model variables can now be parametrized through either of `shape`, `dims` or `size` (see [#4625](https://github.com/pymc-devs/pymc3/pull/4625)): - - With `shape` the length of dimensions must be given numerically or as scalar Aesara `Variables`. A `SpecifyShape` `Op` is added automatically unless `Ellipsis` is used. Using `shape` restricts the model variable to the exact length and re-sizing is no longer possible. + - With `shape` the length of all dimensions must be given numerically or as scalar Aesara `Variables`. A `SpecifyShape` `Op` is added automatically unless `Ellipsis` is used. Using `shape` restricts the model variable to the exact length and re-sizing is no longer possible. - `dims` keeps model variables re-sizeable (for example through `pm.Data`) and leads to well defined coordinates in `InferenceData` objects. - - The `size` kwarg creates new dimensions in addition to what is implied by RV parameters. - - An `Ellipsis` (`...`) in the last position of `shape` or `dims` can be used as short-hand notation for implied dimensions. + - The `size` kwarg resembles the behavior found in Aesara and NumPy: It does not include _support_ dimensions. + - An `Ellipsis` (`...`) in the last position of either kwarg can be used as short-hand notation for implied dimensions. - ... ### Maintenance diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index a6f9922371..be9e5dd628 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -49,9 +49,9 @@ PLATFORM = sys.platform -Shape = Union[int, Sequence[Union[str, type(Ellipsis)]], Variable] -Dims = Union[str, Sequence[Union[str, None, type(Ellipsis)]]] -Size = Union[int, Tuple[int, ...]] +Shape = Union[int, Sequence[Union[int, type(Ellipsis), Variable]], Variable] +Dims = Union[str, Sequence[Union[str, type(Ellipsis), None]]] +Size = Union[int, Sequence[Union[int, type(Ellipsis), Variable]], Variable] class _Unpickling: @@ -139,7 +139,7 @@ def _validate_shape_dims_size( raise ValueError("The `shape` parameter must be an int, list or tuple.") if not isinstance(dims, (type(None), str, list, tuple)): raise ValueError("The `dims` parameter must be a str, list or tuple.") - if not isinstance(size, (type(None), int, list, tuple)): + if not isinstance(size, (type(None), int, list, tuple, Variable)): raise ValueError("The `size` parameter must be an int, list or tuple.") # Auto-convert non-tupled parameters @@ -155,17 +155,14 @@ def _validate_shape_dims_size( shape = tuple(shape) if not isinstance(dims, (type(None), tuple)): dims = tuple(dims) - if not isinstance(size, (type(None), tuple)): + if not isinstance(size, (type(None), tuple, Variable)): size = tuple(size) - if not _valid_ellipsis_position(shape): - raise ValueError( - f"Ellipsis in `shape` may only appear in the last position. Actual: {shape}" - ) - if not _valid_ellipsis_position(dims): - raise ValueError(f"Ellipsis in `dims` may only appear in the last position. Actual: {dims}") - if size is not None and Ellipsis in size: - raise ValueError(f"The `size` parameter cannot contain an Ellipsis. Actual: {size}") + for kwarg, val in dict(shape=shape, dims=dims, size=size).items(): + if not _valid_ellipsis_position(val): + raise ValueError( + f"Ellipsis in `{kwarg}` may only appear in the last position. Actual: {val}" + ) return shape, dims, size @@ -240,12 +237,12 @@ def __new__( rng = model.default_rng _, dims, _ = _validate_shape_dims_size(dims=dims) - resize = None + batch_shape = None # Create the RV without specifying testval, because the testval may have a shape # that only matches after replicating with a size implied by dims (see below). rv_out = cls.dist(*args, rng=rng, testval=None, **kwargs) - n_implied = rv_out.ndim + ndim_implied = rv_out.ndim # The `.dist()` can wrap automatically with a SpecifyShape Op which brings informative # error messages earlier in model construction. @@ -261,33 +258,33 @@ def __new__( # Auto-complete the dims tuple to the full length dims = (*dims[:-1], *[None] * rv_out.ndim) - n_resize = len(dims) - n_implied + ndim_batch = len(dims) - ndim_implied - # All resize dims must be known already (numerically or symbolically). - unknown_resize_dims = set(dims[:n_resize]) - set(model.dim_lengths) - if unknown_resize_dims: + # All batch dims must be known already (numerically or symbolically). + unknown_batch_dims = set(dims[:ndim_batch]) - set(model.dim_lengths) + if unknown_batch_dims: raise KeyError( - f"Dimensions {unknown_resize_dims} are unknown to the model and cannot be used to specify a `size`." + f"Dimensions {unknown_batch_dims} are unknown to the model and cannot be used to specify a `size`." ) - # The numeric/symbolic resize tuple can be created using model.RV_dim_lengths - resize = tuple(model.dim_lengths[dname] for dname in dims[:n_resize]) + # The numeric/symbolic batch shape can be created using model.RV_dim_lengths + batch_shape = tuple(model.dim_lengths[dname] for dname in dims[:ndim_batch]) elif observed is not None: if not hasattr(observed, "shape"): observed = pandas_to_array(observed) - n_resize = observed.ndim - n_implied - resize = tuple(observed.shape[d] for d in range(n_resize)) + ndim_batch = observed.ndim - ndim_implied + batch_shape = tuple(observed.shape[d] for d in range(ndim_batch)) - if resize: + if batch_shape: # A batch size was specified through `dims`, or implied by `observed`. - rv_out = change_rv_size(rv_var=rv_out, new_size=resize, expand=True) + rv_out = change_rv_size(rv_var=rv_out, new_size=batch_shape, expand=True) if dims is not None: # Now that we have a handle on the output RV, we can register named implied dimensions that # were not yet known to the model, such that they can be used for size further downstream. - for di, dname in enumerate(dims[n_resize:]): + for di, dname in enumerate(dims[ndim_batch:]): if not dname in model.dim_lengths: - model.add_coord(dname, values=None, length=rv_out.shape[n_resize + di]) + model.add_coord(dname, values=None, length=rv_out.shape[ndim_batch + di]) if testval is not None: # Assigning the testval earlier causes trouble because the RV may not be created with the final shape already. @@ -329,6 +326,9 @@ def dist( size : int, tuple, Variable, optional A scalar or tuple for replicating the RV in addition to its implied shape/dimensionality. + + Ellipsis (...) may be used in the last position of the tuple, + such that only batch dimensions must be specified. testval : optional Test value to be attached to the output RV. Must match its shape exactly. @@ -344,29 +344,46 @@ def dist( shape, _, size = _validate_shape_dims_size(shape=shape, size=size) assert_shape = None - # Create the RV without specifying size or testval. - # The size will be expanded later (if necessary) and only then the testval fits. + # Create the RV without specifying size or testval, because we don't know + # a-priori if `size` contains a batch shape. + # In the end `testval` must match the final shape, but it is + # not taken into consideration when creating the RV. + # Batch-shape expansion (if necessary) of the RV happens later. rv_native = cls.rv_op(*dist_params, size=None, **kwargs) - if shape is None and size is None: - size = () + # Now we know the implied dimensionality and can figure out the batch shape + batch_shape = () + if size is not None: + if not isinstance(size, Variable) and Ellipsis in size: + batch_shape = size[:-1] + else: + # Parametrization through size does not include support dimensions, + # but may include a batch shape. + ndim_support = rv_native.owner.op.ndim_supp + ndim_inputs = rv_native.ndim - ndim_support + # Be careful to avoid len(size) because it may be symbolic + if ndim_inputs == 0: + batch_shape = size + else: + batch_shape = size[:-ndim_inputs] elif shape is not None: - # SpecifyShape is automatically applied for symbolic and non-Ellipsis shapes - if isinstance(shape, Variable): - assert_shape = shape - size = () + if not isinstance(shape, Variable) and Ellipsis in shape: + # Can't assert a shape without knowing all dimension lengths. + # The batch shape are all entries before ... + batch_shape = tuple(shape[:-1]) else: - if Ellipsis in shape: - size = tuple(shape[:-1]) + # Fully symbolic, or without Ellipsis shapes can be asserted. + assert_shape = shape + # The batch shape are the entries that preceed the implied dimensions. + # Be careful to avoid len(shape) because it may be symbolic + if rv_native.ndim == 0: + batch_shape = shape else: - size = tuple(shape[: len(shape) - rv_native.ndim]) - assert_shape = shape - # no-op conditions: - # `elif size is not None` (User already specified how to expand the RV) - # `else` (Unreachable) - - if size: - rv_out = change_rv_size(rv_var=rv_native, new_size=size, expand=True) + batch_shape = shape[: -rv_native.ndim] + # else: both dimensionality kwargs are None + + if batch_shape: + rv_out = change_rv_size(rv_var=rv_native, new_size=batch_shape, expand=True) else: rv_out = rv_native diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py index 4eb6b01817..c4a608a4e4 100644 --- a/pymc3/distributions/multivariate.py +++ b/pymc3/distributions/multivariate.py @@ -26,7 +26,6 @@ from aesara.graph.op import Op from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal -from aesara.tensor.random.utils import broadcast_params from aesara.tensor.slinalg import ( Cholesky, Solve, @@ -427,27 +426,6 @@ def _distr_parameters_for_repr(self): return ["a"] -class MultinomialRV(MultinomialRV): - """Aesara's `MultinomialRV` doesn't broadcast; this one does.""" - - @classmethod - def rng_fn(cls, rng, n, p, size): - if n.ndim > 0 or p.ndim > 1: - n, p = broadcast_params([n, p], cls.ndims_params) - size = tuple(size or ()) - - if size: - n = np.broadcast_to(n, size + n.shape) - p = np.broadcast_to(p, size + p.shape) - - res = np.empty(p.shape) - for idx in np.ndindex(p.shape[:-1]): - res[idx] = rng.multinomial(n[idx], p[idx]) - return res - else: - return rng.multinomial(n, p, size=size) - - multinomial = MultinomialRV() diff --git a/pymc3/tests/test_data_container.py b/pymc3/tests/test_data_container.py index 1ca14088cc..c58619e8f4 100644 --- a/pymc3/tests/test_data_container.py +++ b/pymc3/tests/test_data_container.py @@ -163,7 +163,7 @@ def test_shared_data_as_rv_input(self): with pm.Model() as m: x = pm.Data("x", [1.0, 2.0, 3.0]) assert x.eval().shape == (3,) - y = pm.Normal("y", mu=x, size=2) + y = pm.Normal("y", mu=x, shape=(2, ...)) assert y.eval().shape == (2, 3) idata = pm.sample( chains=1, diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py index 63466b7902..9013918575 100644 --- a/pymc3/tests/test_distributions.py +++ b/pymc3/tests/test_distributions.py @@ -2074,29 +2074,33 @@ def test_multinomial_mode(self, p, n): @pytest.mark.parametrize( "p, size, n", [ - [[0.25, 0.25, 0.25, 0.25], (4,), 2], - [[0.25, 0.25, 0.25, 0.25], (1, 4), 3], + [[0.25, 0.25, 0.25, 0.25], (7,), 2], + [[0.25, 0.25, 0.25, 0.25], (1, 7), 3], # 3: expect to fail - # [[.25, .25, .25, .25], (10, 4)], - [[0.25, 0.25, 0.25, 0.25], (10, 1, 4), 5], + # [[.25, .25, .25, .25], (10, 7)], + [[0.25, 0.25, 0.25, 0.25], (10, 1, 7), 5], # 5: expect to fail - # [[[.25, .25, .25, .25]], (2, 4), [7, 11]], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (2, 4), 13], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (1, 2, 4), [23, 29]], + # [[[.25, .25, .25, .25]], (2, 5), [7, 11]], + [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (5, 2), 13], + [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (5, 7, 2), [23, 29]], [ [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], - (10, 2, 4), + (10, 8, 2), [31, 37], ], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (2, 4), [17, 19]], + [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (3, 2), [17, 19]], ], ) def test_multinomial_random(self, p, size, n): p = np.asarray(p) with Model() as model: m = Multinomial("m", n=n, p=p, size=size) - - assert m.eval().shape == size + p.shape + # The support has length 4 in all test parametrizations! + # Broadcasting of the `p` parameter does not affect the ndim_supp + # of the Op, hence the broadcasted p must be included in `size`. + support_shape = (p.shape[-1],) + assert support_shape == (4,) + assert m.eval().shape == size + support_shape @pytest.mark.skip(reason="Moment calculations have not been refactored yet") def test_multinomial_mode_with_shape(self): @@ -2197,7 +2201,7 @@ def test_batch_multinomial(self): decimal=select_by_precision(float64=6, float32=3), ) - dist = Multinomial.dist(n=n, p=p, size=2) + dist = Multinomial.dist(n=n, p=p, size=(2, ...)) sample = dist.eval() assert_allclose(sample, np.stack([vals, vals], axis=0)) diff --git a/pymc3/tests/test_model.py b/pymc3/tests/test_model.py index f307ad93c5..3d8a478ac8 100644 --- a/pymc3/tests/test_model.py +++ b/pymc3/tests/test_model.py @@ -35,6 +35,7 @@ from pymc3 import Deterministic, Potential from pymc3.blocking import DictToArrayBijection, RaveledVars from pymc3.distributions import Normal, logpt_sum, transforms +from pymc3.exceptions import ShapeError from pymc3.model import Point, ValueGradFunction from pymc3.tests.helpers import SeededTest @@ -465,6 +466,10 @@ def test_make_obs_var(): # Create the testval attribute simply for the sake of model testing fake_distribution.name = input_name + # The function requires data and RV dimensionality to be compatible + with pytest.raises(ShapeError, match="Dimensionality of data and RV don't match."): + fake_model.make_obs_var(fake_distribution, np.ones((3, 3, 1)), None, None) + # Check function behavior using the various inputs # dense, sparse: Ensure that the missing values are appropriately set to None # masked: a deterministic variable is returned diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py index 2581c0ae47..caf6f0079d 100644 --- a/pymc3/tests/test_sampling.py +++ b/pymc3/tests/test_sampling.py @@ -967,14 +967,14 @@ def test_multivariate2(self): def test_layers(self): with pm.Model() as model: - a = pm.Uniform("a", lower=0, upper=1, size=5) - b = pm.Binomial("b", n=1, p=a, size=7) + a = pm.Uniform("a", lower=0, upper=1, size=10) + b = pm.Binomial("b", n=1, p=a) model.default_rng.get_value(borrow=True).seed(232093) b_sampler = aesara.function([], b) avg = np.stack([b_sampler() for i in range(10000)]).mean(0) - npt.assert_array_almost_equal(avg, 0.5 * np.ones((7, 5)), decimal=2) + npt.assert_array_almost_equal(avg, 0.5 * np.ones((10,)), decimal=2) def test_transformed(self): n = 18 diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index a64d1de5bf..23adcd1dba 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -18,6 +18,7 @@ import pytest from aesara import tensor as at +from aesara.tensor.shape import SpecifyShape import pymc3 as pm @@ -225,70 +226,81 @@ def test_sample_generate_values(fixture_model, fixture_sizes): class TestShapeDimsSize: - @pytest.mark.parametrize("param_shape", [(), (3,)]) - @pytest.mark.parametrize("batch_shape", [(), (3,)]) - @pytest.mark.parametrize( - "parametrization", - [ - "implicit", - "shape", - "shape...", - "dims", - "dims...", - "size", - ], - ) - def test_param_and_batch_shape_combos( - self, param_shape: tuple, batch_shape: tuple, parametrization: str - ): + @pytest.mark.parametrize("support_shape", [(), (9,)]) + @pytest.mark.parametrize("input_shape", [(), (2,), (3, 5)]) + @pytest.mark.parametrize("batch_shape", [(), (6,), (7, 8)]) + def test_parametrization_combos(self, support_shape, input_shape, batch_shape): + ndim_batch = len(batch_shape) + ndim_inputs = len(input_shape) + ndim_support = len(support_shape) + ndim = ndim_batch + ndim_inputs + ndim_support + + if ndim_support == 0: + dist = pm.Normal + inputs = dict(mu=np.ones(input_shape)) + expected = batch_shape + input_shape + elif ndim_support == 1: + dist = pm.MvNormal + mu_shape = input_shape + support_shape + inputs = dict(mu=np.ones(mu_shape), cov=np.eye(support_shape[0])) + expected = batch_shape + input_shape + support_shape + else: + raise NotImplementedError( + f"No tests implemented for {ndim_support}-dimensional RV support." + ) + + # Without dimensionality kwargs, the RV shape depends only on its inputs and support + assert dist.dist(**inputs).eval().shape == input_shape + support_shape + + # The `shape` includes the support dims (0 in the case of univariates). + assert ( + dist.dist(**inputs, shape=(*batch_shape, *input_shape, *support_shape)).eval().shape + == expected + ) + + # In contrast, `size` is without support dims + assert dist.dist(**inputs, size=(*batch_shape, *input_shape)).eval().shape == expected + + # With Ellipsis in the last position, `shape` is independent of all parameter and support dims. + assert dist.dist(**inputs, shape=(*batch_shape, ...)).eval().shape == expected + + # This test uses fixed-length dimensions that are specified through model coords. + # Here those coords are created depending on the test parametrization. coords = {} - param_dims = [] + support_dims = [] + input_dims = [] batch_dims = [] - # Create coordinates corresponding to the parameter shape - for d in param_shape: - dname = f"param_dim_{d}" + for d in support_shape: + dname = f"support_dim_{d}" coords[dname] = [f"c_{i}" for i in range(d)] - param_dims.append(dname) - assert len(param_dims) == len(param_shape) - # Create coordinates corresponding to the batch shape + support_dims.append(dname) + assert len(support_dims) == ndim_support + + for d in input_shape: + dname = f"input_dim_{d}" + coords[dname] = [f"c_{i}" for i in range(d)] + input_dims.append(dname) + assert len(input_dims) == ndim_inputs + for d in batch_shape: dname = f"batch_dim_{d}" coords[dname] = [f"c_{i}" for i in range(d)] batch_dims.append(dname) - assert len(batch_dims) == len(batch_shape) + assert len(batch_dims) == ndim_batch + # The `dims` are only available with a model. with pm.Model(coords=coords) as pmodel: - mu = aesara.shared(np.random.normal(size=param_shape)) + rv_dims_full = dist( + "rv_dims_full", **inputs, dims=batch_dims + input_dims + support_dims + ) + assert rv_dims_full.eval().shape == expected + assert len(pmodel.RV_dims["rv_dims_full"]) == ndim + assert rv_dims_full.eval().shape == expected - with pytest.warns(None): - if parametrization == "implicit": - rv = pm.Normal("rv", mu=mu).shape == param_shape - else: - if parametrization == "shape": - rv = pm.Normal("rv", mu=mu, shape=batch_shape + param_shape) - assert rv.eval().shape == batch_shape + param_shape - elif parametrization == "shape...": - rv = pm.Normal("rv", mu=mu, shape=(*batch_shape, ...)) - assert rv.eval().shape == batch_shape + param_shape - elif parametrization == "dims": - rv = pm.Normal("rv", mu=mu, dims=batch_dims + param_dims) - assert rv.eval().shape == batch_shape + param_shape - elif parametrization == "dims...": - rv = pm.Normal("rv", mu=mu, dims=(*batch_dims, ...)) - n_size = len(batch_shape) - n_implied = len(param_shape) - ndim = n_size + n_implied - assert len(pmodel.RV_dims["rv"]) == ndim, pmodel.RV_dims - assert len(pmodel.RV_dims["rv"][:n_size]) == len(batch_dims) - assert len(pmodel.RV_dims["rv"][n_size:]) == len(param_dims) - if n_implied > 0: - assert pmodel.RV_dims["rv"][-1] is None - elif parametrization == "size": - rv = pm.Normal("rv", mu=mu, size=batch_shape) - assert rv.eval().shape == batch_shape + param_shape - else: - raise NotImplementedError("Invalid test case parametrization.") + rv_dims_ellipsis = dist("rv_dims_ellipsis", **inputs, dims=(*batch_dims, ...)) + assert len(pmodel.RV_dims["rv_dims_ellipsis"]) == ndim + assert rv_dims_ellipsis.eval().shape == expected def test_define_dims_on_the_fly(self): with pm.Model() as pmodel: @@ -316,7 +328,7 @@ def test_data_defined_size_dimension_can_register_dimname(self): assert "first" in pmodel.dim_lengths assert "second" in pmodel.dim_lengths # two dimensions are implied; a "third" dimension is created - y = pm.Normal("y", mu=x, size=2, dims=("third", "first", "second")) + y = pm.Normal("y", mu=x, shape=(2, ...), dims=("third", "first", "second")) assert "third" in pmodel.dim_lengths assert y.eval().shape() == (2, 1, 4) @@ -346,20 +358,48 @@ def test_observed_with_column_vector(self): pm.Normal("x2", mu=0, sd=1, observed=np.random.normal(size=(3, 1))) model.logp() - def test_dist_api_works(self): + def test_dist_api_basics(self): mu = aesara.shared(np.array([1, 2, 3])) with pytest.raises(NotImplementedError, match="API is not yet supported"): pm.Normal.dist(mu=mu, dims=("town",)) assert pm.Normal.dist(mu=mu, shape=(3,)).eval().shape == (3,) assert pm.Normal.dist(mu=mu, shape=(5, 3)).eval().shape == (5, 3) assert pm.Normal.dist(mu=mu, shape=(7, ...)).eval().shape == (7, 3) - assert pm.Normal.dist(mu=mu, size=(4,)).eval().shape == (4, 3) + assert pm.Normal.dist(mu=mu, size=(4, 3)).eval().shape == (4, 3) + assert pm.Normal.dist(mu=mu, size=(8, ...)).eval().shape == (8, 3) + + def test_tensor_shape(self): + s1 = at.scalar(dtype="int32") + rv = pm.Uniform.dist(1, [1, 2], shape=[s1, 2]) + f = aesara.function([s1], rv, mode=aesara.Mode("py")) + assert f(3).shape == (3, 2) + assert f(7).shape == (7, 2) + + # As long as its length is fixed, it can also be a vector Variable + rv_a = pm.Normal.dist(size=(7, 3)) + assert rv_a.shape.ndim == 1 + rv_b = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), shape=rv_a.shape) + assert rv_b.eval().shape == (7, 3) + + def test_tensor_size(self): + s1 = at.scalar(dtype="int32") + s2 = at.scalar(dtype="int32") + rv = pm.Uniform.dist(1, [1, 2], size=[s1, s2, ...]) + f = aesara.function([s1, s2], rv, mode=aesara.Mode("py")) + assert f(3, 5).shape == (3, 5, 2) + assert f(7, 3).shape == (7, 3, 2) + + # As long as its length is fixed, it can also be a vector Variable + rv_a = pm.Normal.dist(size=(7, 4)) + assert rv_a.shape.ndim == 1 + rv_b = pm.MvNormal.dist(mu=np.ones(3), cov=np.eye(3), size=rv_a.shape) + assert rv_b.eval().shape == (7, 4, 3) def test_auto_assert_shape(self): with pytest.raises(AssertionError, match="will never match"): pm.Normal.dist(mu=[1, 2], shape=[]) - mu = at.vector(name="mu_input") + mu = at.vector() rv = pm.Normal.dist(mu=mu, shape=[3, 4]) f = aesara.function([mu], rv, mode=aesara.Mode("py")) assert f([1, 2, 3, 4]).shape == (3, 4) @@ -368,24 +408,20 @@ def test_auto_assert_shape(self): f([1, 2]) # The `shape` can be symbolic! - s = at.vector(dtype="int32") - rv = pm.Uniform.dist(2, [4, 5], shape=s) - f = aesara.function([s], rv, mode=aesara.Mode("py")) - f( - [ - 2, - ] - ) + # This example has a batch dimension too, so under the hood it + # becomes a symbolic input to `specify_shape` AND a symbolic `batch_shape`. + s1 = at.scalar(dtype="int32") + s2 = at.scalar(dtype="int32") + rv = pm.Uniform.dist(2, [4, 5], shape=[s1, s2]) + assert isinstance(rv.owner.op, SpecifyShape) + f = aesara.function([s1, s2], rv, mode=aesara.Mode("py")) + assert f(3, 2).shape == (3, 2) + assert f(7, 2).shape == (7, 2) with pytest.raises( AssertionError, - match=r"Got 1 dimensions \(shape \(2,\)\), expected 2 dimensions with shape \(3, 4\).", + match=r"Got shape \(3, 2\), expected \(3, 4\).", ): - f([3, 4]) - with pytest.raises( - AssertionError, - match=r"Got 1 dimensions \(shape \(2,\)\), expected 0 dimensions with shape \(\).", - ): - f([]) + f(3, 4) pass def test_lazy_flavors(self): @@ -394,7 +430,7 @@ def test_lazy_flavors(self): _validate_shape_dims_size(dims="town") _validate_shape_dims_size(size=7) - assert pm.Uniform.dist(2, [4, 5], size=[3, 4]).eval().shape == (3, 4, 2) + assert pm.Uniform.dist(2, [4, 5], size=[3, 4, 2]).eval().shape == (3, 4, 2) assert pm.Uniform.dist(2, [4, 5], shape=[3, 2]).eval().shape == (3, 2) with pm.Model(coords=dict(town=["Greifswald", "Madrid"])): assert pm.Normal("n2", mu=[1, 2], dims=("town",)).eval().shape == (2,) @@ -417,5 +453,5 @@ def test_invalid_flavors(self): _validate_shape_dims_size(shape=(3, ..., 2)) with pytest.raises(ValueError, match="may only appear in the last position"): _validate_shape_dims_size(dims=(..., "town")) - with pytest.raises(ValueError, match="cannot contain"): - _validate_shape_dims_size(size=(3, ...)) + with pytest.raises(ValueError, match="may only appear in the last position"): + _validate_shape_dims_size(size=(3, ..., ...)) diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py index 6ef93eb5d1..20a4f058c0 100644 --- a/pymc3/tests/test_step.py +++ b/pymc3/tests/test_step.py @@ -982,7 +982,7 @@ def test_linalg(self, caplog): a = Normal("a", size=2, testval=floatX(np.zeros(2))) a = at.switch(a > 0, np.inf, a) b = at.slinalg.solve(floatX(np.eye(2)), a) - Normal("c", mu=b, shape=(2,), testval=floatX(np.r_[0.0, 0.0])) + Normal("c", mu=b, size=(2,), testval=floatX(np.r_[0.0, 0.0])) caplog.clear() trace = sample(20, init=None, tune=5, chains=2) warns = [msg.msg for msg in caplog.records] diff --git a/pymc3/tests/test_transforms.py b/pymc3/tests/test_transforms.py index c2a8f0eb00..a165c3b13d 100644 --- a/pymc3/tests/test_transforms.py +++ b/pymc3/tests/test_transforms.py @@ -274,11 +274,11 @@ def test_chain_jacob_det(): class TestElementWiseLogp(SeededTest): - def build_model(self, distfam, params, *, size=None, shape=None, transform=None, testval=None): + def build_model(self, distfam, params, *, size=None, transform=None, testval=None): if testval is not None: testval = pm.floatX(testval) with pm.Model() as m: - distfam("x", size=size, shape=shape, transform=transform, testval=testval, **params) + distfam("x", size=size, transform=transform, testval=testval, **params) return m def check_transform_elementwise_logp(self, model): @@ -334,28 +334,26 @@ def test_exponential(self, lam, size): self.check_transform_elementwise_logp(model) @pytest.mark.parametrize( - "a,b,shape", + "a,b,size", [ (1.0, 1.0, 2), (0.5, 0.5, (2, 3)), (np.ones(3), np.ones(3), (4, 3)), ], ) - def test_beta(self, a, b, shape): - model = self.build_model( - pm.Beta, {"alpha": a, "beta": b}, shape=shape, transform=tr.logodds - ) + def test_beta(self, a, b, size): + model = self.build_model(pm.Beta, {"alpha": a, "beta": b}, size=size, transform=tr.logodds) self.check_transform_elementwise_logp(model) @pytest.mark.parametrize( - "lower,upper,shape", + "lower,upper,size", [ (0.0, 1.0, 2), (0.5, 5.5, (2, 3)), (pm.floatX(np.zeros(3)), pm.floatX(np.ones(3)), (4, 3)), ], ) - def test_uniform(self, lower, upper, shape): + def test_uniform(self, lower, upper, size): def transform_params(rv_var): _, _, _, lower, upper = rv_var.owner.inputs lower = at.as_tensor_variable(lower) if lower is not None else None @@ -364,7 +362,7 @@ def transform_params(rv_var): interval = tr.Interval(transform_params) model = self.build_model( - pm.Uniform, {"lower": lower, "upper": upper}, shape=shape, transform=interval + pm.Uniform, {"lower": lower, "upper": upper}, size=size, transform=interval ) self.check_transform_elementwise_logp(model) @@ -394,15 +392,15 @@ def transform_params(rv_var): ) def test_vonmises(self, mu, kappa, shape): model = self.build_model( - pm.VonMises, {"mu": mu, "kappa": kappa}, shape=shape, transform=tr.circular + pm.VonMises, {"mu": mu, "kappa": kappa}, size=size, transform=tr.circular ) self.check_transform_elementwise_logp(model) @pytest.mark.parametrize( - "a,shape", [(np.ones(2), None), (np.ones((2, 3)) * 0.5, None), (np.ones(3), (4,))] + "a,size", [(np.ones(2), None), (np.ones((2, 3)) * 0.5, None), (np.ones(3), (4,))] ) - def test_dirichlet(self, a, shape): - model = self.build_model(pm.Dirichlet, {"a": a}, shape=shape, transform=tr.stick_breaking) + def test_dirichlet(self, a, size): + model = self.build_model(pm.Dirichlet, {"a": a}, size=size, transform=tr.stick_breaking) self.check_vectortransform_elementwise_logp(model, vect_opt=1) def test_normal_ordered(self): @@ -416,59 +414,59 @@ def test_normal_ordered(self): self.check_vectortransform_elementwise_logp(model, vect_opt=0) @pytest.mark.parametrize( - "sd,shape", + "sd,size", [ (2.5, (2,)), (np.ones(3), (4, 3)), ], ) @pytest.mark.xfail(condition=(aesara.config.floatX == "float32"), reason="Fails on float32") - def test_half_normal_ordered(self, sd, shape): - testval = np.sort(np.abs(np.random.randn(*shape))) + def test_half_normal_ordered(self, sd, size): + testval = np.sort(np.abs(np.random.randn(*size))) model = self.build_model( pm.HalfNormal, {"sd": sd}, - shape=shape, + size=size, testval=testval, transform=tr.Chain([tr.log, tr.ordered]), ) self.check_vectortransform_elementwise_logp(model, vect_opt=0) - @pytest.mark.parametrize("lam,shape", [(2.5, (2,)), (np.ones(3), (4, 3))]) - def test_exponential_ordered(self, lam, shape): - testval = np.sort(np.abs(np.random.randn(*shape))) + @pytest.mark.parametrize("lam,size", [(2.5, (2,)), (np.ones(3), (4, 3))]) + def test_exponential_ordered(self, lam, size): + testval = np.sort(np.abs(np.random.randn(*size))) model = self.build_model( pm.Exponential, {"lam": lam}, - shape=shape, + size=size, testval=testval, transform=tr.Chain([tr.log, tr.ordered]), ) self.check_vectortransform_elementwise_logp(model, vect_opt=0) @pytest.mark.parametrize( - "a,b,shape", + "a,b,size", [ (1.0, 1.0, (2,)), (np.ones(3), np.ones(3), (4, 3)), ], ) - def test_beta_ordered(self, a, b, shape): - testval = np.sort(np.abs(np.random.rand(*shape))) + def test_beta_ordered(self, a, b, size): + testval = np.sort(np.abs(np.random.rand(*size))) model = self.build_model( pm.Beta, {"alpha": a, "beta": b}, - shape=shape, + size=size, testval=testval, transform=tr.Chain([tr.logodds, tr.ordered]), ) self.check_vectortransform_elementwise_logp(model, vect_opt=0) @pytest.mark.parametrize( - "lower,upper,shape", + "lower,upper,size", [(0.0, 1.0, (2,)), (pm.floatX(np.zeros(3)), pm.floatX(np.ones(3)), (4, 3))], ) - def test_uniform_ordered(self, lower, upper, shape): + def test_uniform_ordered(self, lower, upper, size): def transform_params(rv_var): _, _, _, lower, upper = rv_var.owner.inputs lower = at.as_tensor_variable(lower) if lower is not None else None @@ -477,11 +475,11 @@ def transform_params(rv_var): interval = tr.Interval(transform_params) - testval = np.sort(np.abs(np.random.rand(*shape))) + testval = np.sort(np.abs(np.random.rand(*size))) model = self.build_model( pm.Uniform, {"lower": lower, "upper": upper}, - shape=shape, + size=size, testval=testval, transform=tr.Chain([interval, tr.ordered]), ) @@ -495,26 +493,26 @@ def test_vonmises_ordered(self, mu, kappa, shape): model = self.build_model( pm.VonMises, {"mu": mu, "kappa": kappa}, - shape=shape, + size=size, testval=testval, transform=tr.Chain([tr.circular, tr.ordered]), ) self.check_vectortransform_elementwise_logp(model, vect_opt=0) @pytest.mark.parametrize( - "lower,upper,shape,transform", + "lower,upper,size,transform", [ (0.0, 1.0, (2,), tr.stick_breaking), (0.5, 5.5, (2, 3), tr.stick_breaking), (np.zeros(3), np.ones(3), (4, 3), tr.Chain([tr.sum_to_1, tr.logodds])), ], ) - def test_uniform_other(self, lower, upper, shape, transform): - testval = np.ones(shape) / shape[-1] + def test_uniform_other(self, lower, upper, size, transform): + testval = np.ones(size) / size[-1] model = self.build_model( pm.Uniform, {"lower": lower, "upper": upper}, - shape=shape, + size=size, testval=testval, transform=transform, )