From 5007d9739703718be1bcfd012bd4b7ceaa2cbb05 Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Mon, 19 Apr 2021 11:33:36 +0200 Subject: [PATCH] Automatically add SpecifyShape Op when full-length shape is given --- RELEASE-NOTES.md | 2 +- pymc3/aesaraf.py | 3 +++ pymc3/distributions/distribution.py | 28 +++++++++++++++++++++++- pymc3/tests/test_logp.py | 2 +- pymc3/tests/test_shape_handling.py | 33 +++++++++++++++++++++++++++++ 5 files changed, 65 insertions(+), 3 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 49d013f3c0..b7918cb30c 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -9,7 +9,7 @@ ### 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`. Using `shape` restricts the model variable to the exact length and re-sizing is no longer possible. + - 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. - `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. diff --git a/pymc3/aesaraf.py b/pymc3/aesaraf.py index e30248b841..5a99849fd2 100644 --- a/pymc3/aesaraf.py +++ b/pymc3/aesaraf.py @@ -45,6 +45,7 @@ from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream from aesara.tensor.elemwise import Elemwise from aesara.tensor.random.op import RandomVariable +from aesara.tensor.shape import SpecifyShape from aesara.tensor.sharedvar import SharedVariable from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1 from aesara.tensor.var import TensorVariable @@ -146,6 +147,8 @@ def change_rv_size( Expand the existing size by `new_size`. """ + if isinstance(rv_var.owner.op, SpecifyShape): + rv_var = rv_var.owner.inputs[0] rv_node = rv_var.owner rng, size, dtype, *dist_params = rv_node.inputs name = rv_var.name diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index 8a0c43b07b..3fba3cf7a2 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -28,6 +28,7 @@ from aesara.graph.basic import Variable from aesara.tensor.random.op import RandomVariable +from aesara.tensor.shape import SpecifyShape, specify_shape from pymc3.aesaraf import change_rv_size, pandas_to_array from pymc3.distributions import _logcdf, _logp @@ -253,6 +254,13 @@ def __new__( rv_out = cls.dist(*args, rng=rng, testval=None, **kwargs) n_implied = rv_out.ndim + # The `.dist()` can wrap automatically with a SpecifyShape Op which brings informative + # error messages earlier in model construction. + # Here, however, the underyling RV must be used - a new SpecifyShape Op can be added at the end. + assert_shape = None + if isinstance(rv_out.owner.op, SpecifyShape): + rv_out, assert_shape = rv_out.owner.inputs + # `dims` are only available with this API, because `.dist()` can be used # without a modelcontext and dims are not tracked at the Aesara level. if dims is not None: @@ -292,7 +300,15 @@ def __new__( # Assigning the testval earlier causes trouble because the RV may not be created with the final shape already. rv_out.tag.test_value = testval - return model.register_rv(rv_out, name, observed, total_size, dims=dims, transform=transform) + rv_registered = model.register_rv( + rv_out, name, observed, total_size, dims=dims, transform=transform + ) + + # Wrapping in specify_shape now does not break transforms: + if assert_shape is not None: + rv_registered = specify_shape(rv_registered, assert_shape) + + return rv_registered @classmethod def dist( @@ -314,6 +330,9 @@ def dist( Ellipsis (...) may be used in the last position of the tuple, and automatically expand to the shape implied by RV inputs. + + Without Ellipsis, a `SpecifyShape` Op is automatically applied, + constraining this model variable to exactly the specified shape. size : int, tuple, Variable, optional A scalar or tuple for replicating the RV in addition to its implied shape/dimensionality. @@ -330,6 +349,7 @@ def dist( raise NotImplementedError("The use of a `.dist(dims=...)` API is not yet supported.") 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. @@ -338,13 +358,16 @@ def dist( if shape is None and size is None: size = () elif shape is not None: + # SpecifyShape is automatically applied for symbolic and non-Ellipsis shapes if isinstance(shape, Variable): + assert_shape = shape size = () else: if Ellipsis in shape: size = tuple(shape[:-1]) 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) @@ -354,6 +377,9 @@ def dist( else: rv_out = rv_native + if assert_shape is not None: + rv_out = specify_shape(rv_out, shape=assert_shape) + if testval is not None: rv_out.tag.test_value = testval diff --git a/pymc3/tests/test_logp.py b/pymc3/tests/test_logp.py index 215e155e2f..f53c640a8f 100644 --- a/pymc3/tests/test_logp.py +++ b/pymc3/tests/test_logp.py @@ -86,7 +86,7 @@ def test_logpt_incsubtensor(indices, shape): sigma = 0.001 rng = aesara.shared(np.random.RandomState(232), borrow=True) - a = Normal.dist(mu, sigma, shape=shape, rng=rng) + a = Normal.dist(mu, sigma, rng=rng) a.name = "a" a_idx = at.set_subtensor(a[indices], data) diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index d181365ada..3c93721f15 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -350,6 +350,39 @@ def test_dist_api_works(self): assert pm.Normal.dist(mu=mu, shape=(7, ...)).eval().shape == (7, 3) assert pm.Normal.dist(mu=mu, size=(4,)).eval().shape == (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") + 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) + + with pytest.raises(AssertionError, match=r"Got shape \(3, 2\), expected \(3, 4\)."): + 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, + ] + ) + with pytest.raises( + AssertionError, + match=r"Got 1 dimensions \(shape \(2,\)\), expected 2 dimensions with shape \(3, 4\).", + ): + f([3, 4]) + with pytest.raises( + AssertionError, + match=r"Got 1 dimensions \(shape \(2,\)\), expected 0 dimensions with shape \(\).", + ): + f([]) + pass + def test_lazy_flavors(self): _validate_shape_dims_size(shape=5)