From d4894a7b5a3ee720ed56361e3d4e939a75660d07 Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Sat, 24 Apr 2021 18:06:24 +0200 Subject: [PATCH] Revert functionality change of size kwarg Corresponding tests were reverted, or edited to use other parametrization flavors. The Ellipsis feature now works with all three dimensionality kwargs. The MultinomialRV implementation was removed, because the broadcasting behavior was implemented in Aesara. Closes #4662 --- RELEASE-NOTES.md | 6 +- pymc3/distributions/distribution.py | 109 +++++++++------- pymc3/distributions/multivariate.py | 22 ---- pymc3/tests/test_data_container.py | 2 +- pymc3/tests/test_distributions.py | 28 +++-- pymc3/tests/test_model.py | 5 + pymc3/tests/test_sampling.py | 6 +- pymc3/tests/test_shape_handling.py | 186 +++++++++++++++++----------- pymc3/tests/test_step.py | 2 +- pymc3/tests/test_transforms.py | 80 ++++++------ 10 files changed, 241 insertions(+), 205 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index b7918cb30c..3fad1a54f0 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -9,10 +9,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 59d2e185a1..efc622fc1d 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: @@ -146,7 +146,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 @@ -162,17 +162,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 @@ -247,12 +244,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. @@ -268,33 +265,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. @@ -336,6 +333,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. @@ -351,29 +351,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 c6869ec110..0e2949e02a 100644 --- a/pymc3/tests/test_distributions.py +++ b/pymc3/tests/test_distributions.py @@ -1986,29 +1986,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): @@ -2109,7 +2113,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 a155702ba1..3f4ad6322c 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 217cf96010..2a91bd65d9 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 e040e6e244..14cff5e0f3 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,25 +362,25 @@ 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) @pytest.mark.parametrize( - "mu,kappa,shape", [(0.0, 1.0, 2), (-0.5, 5.5, (2, 3)), (np.zeros(3), np.ones(3), (4, 3))] + "mu,kappa,size", [(0.0, 1.0, 2), (-0.5, 5.5, (2, 3)), (np.zeros(3), np.ones(3), (4, 3))] ) @pytest.mark.xfail(reason="Distribution not refactored yet") - def test_vonmises(self, mu, kappa, shape): + def test_vonmises(self, mu, kappa, size): 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): @@ -396,59 +394,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 @@ -457,45 +455,43 @@ 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]), ) self.check_vectortransform_elementwise_logp(model, vect_opt=1) - @pytest.mark.parametrize( - "mu,kappa,shape", [(0.0, 1.0, (2,)), (np.zeros(3), np.ones(3), (4, 3))] - ) + @pytest.mark.parametrize("mu,kappa,size", [(0.0, 1.0, (2,)), (np.zeros(3), np.ones(3), (4, 3))]) @pytest.mark.xfail(reason="Distribution not refactored yet") - def test_vonmises_ordered(self, mu, kappa, shape): - testval = np.sort(np.abs(np.random.rand(*shape))) + def test_vonmises_ordered(self, mu, kappa, size): + testval = np.sort(np.abs(np.random.rand(*size))) 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, )