From c3e82108694c190170b7aa087beed51709fbfe94 Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Fri, 14 May 2021 21:21:42 +0200 Subject: [PATCH 1/9] Disallow "__sample__" as a dimension name Co-authored-by: Thomas Wiecki --- pymc3/model.py | 9 +++++---- pymc3/tests/test_sampling.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pymc3/model.py b/pymc3/model.py index 09018bb22dd..cc80df504ab 100644 --- a/pymc3/model.py +++ b/pymc3/model.py @@ -1000,7 +1000,7 @@ def add_coord( ---------- name : str Name of the dimension. - Forbidden: {"chain", "draw"} + Forbidden: {"chain", "draw", "__sample__"} values : optional, array-like Coordinate values or ``None`` (for auto-numbering). If ``None`` is passed, a ``length`` must be specified. @@ -1008,9 +1008,10 @@ def add_coord( A symbolic scalar of the dimensions length. Defaults to ``aesara.shared(len(values))``. """ - if name in {"draw", "chain"}: + if name in {"draw", "chain", "__sample__"}: raise ValueError( - "Dimensions can not be named `draw` or `chain`, as they are reserved for the sampler's outputs." + "Dimensions can not be named `draw`, `chain` or `__sample__`, " + "as those are reserved for use in `InferenceData`." ) if values is None and length is None: raise ValueError( @@ -1022,7 +1023,7 @@ def add_coord( ) if name in self.coords: if not values.equals(self.coords[name]): - raise ValueError("Duplicate and incompatiple coordinate: %s." % name) + raise ValueError(f"Duplicate and incompatible coordinate: {name}.") else: self._coords[name] = values self._dim_lengths[name] = length or aesara.shared(len(values)) diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py index 9b0a39602a8..f1405bc12ef 100644 --- a/pymc3/tests/test_sampling.py +++ b/pymc3/tests/test_sampling.py @@ -214,7 +214,7 @@ def test_return_inferencedata(self, monkeypatch): return_inferencedata=True, discard_tuned_samples=True, idata_kwargs={"prior": prior}, - random_seed=-1 + random_seed=-1, ) assert "prior" in result assert isinstance(result, InferenceData) From 35bb41a9dc209f5fc792df0a84b739049243cb29 Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Fri, 14 May 2021 21:23:27 +0200 Subject: [PATCH 2/9] Add OS-independent #4652 regression tests --- pymc3/tests/test_ode.py | 4 ++++ pymc3/tests/test_shape_handling.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/pymc3/tests/test_ode.py b/pymc3/tests/test_ode.py index 94dfb0dd6f5..70b935fea42 100644 --- a/pymc3/tests/test_ode.py +++ b/pymc3/tests/test_ode.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import sys import aesara import numpy as np @@ -168,6 +169,9 @@ def ode_func_5(y, t, p): np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._sens_ic) +@pytest.mark.xfail( + condition=sys.platform == "win32", reason="https://github.com/pymc-devs/aesara/issues/390" +) def test_logp_scalar_ode(): """Test the computation of the log probability for these models""" diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index 37c06193226..141c28dad3f 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import aesara import numpy as np import pytest @@ -219,3 +220,32 @@ def test_sample_generate_values(fixture_model, fixture_sizes): prior = pm.sample_prior_predictive(samples=fixture_sizes) for rv in RVs: assert prior[rv.name].shape == size + tuple(rv.distribution.shape) + + +@pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") +def test_size32_doesnt_break_broadcasting(): + size32 = at.constant([1, 10], dtype="int32") + rv = pm.Normal.dist(0, 1, size=size32) + assert rv.broadcastable == (True, False) + + +def test_observed_with_column_vector(): + """This test is related to https://github.com/pymc-devs/aesara/issues/390 which breaks + broadcastability of column-vector RVs. This unexpected change in type can lead to + incompatibilities during graph rewriting for model.logp evaluation. + """ + with pm.Model() as model: + # The `observed` is a broadcastable column vector + obs = at.as_tensor_variable(np.ones((3, 1), dtype=aesara.config.floatX)) + assert obs.broadcastable == (False, True) + + # Both shapes describe broadcastable volumn vectors + size64 = at.constant([3, 1], dtype="int64") + # But the second shape is upcasted from an int32 vector + cast64 = at.cast(at.constant([3, 1], dtype="int32"), dtype="int64") + + pm.Normal("x_size64", mu=0, sd=1, size=size64, observed=obs) + model.logp() + + pm.Normal("x_cast64", mu=0, sd=1, size=cast64, observed=obs) + model.logp() From dbb1f20374eebdda50940d3b0a18de5cf3c73d5f Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Fri, 14 May 2021 21:25:21 +0200 Subject: [PATCH 3/9] Make change_rv_size more robust Co-authored-by: Ricardo Co-authored-by: Thomas Wiecki --- pymc3/aesaraf.py | 18 +++++++++++++++--- pymc3/tests/test_aesaraf.py | 6 ++++++ pymc3/tests/test_distributions_random.py | 9 ++++----- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/pymc3/aesaraf.py b/pymc3/aesaraf.py index bb7becaab80..b61d9b0aa7c 100644 --- a/pymc3/aesaraf.py +++ b/pymc3/aesaraf.py @@ -46,10 +46,12 @@ 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 +from pymc3.exceptions import ShapeError from pymc3.vartypes import continuous_types, int_types, isgenerator, typefilter PotentialShapeType = Union[ @@ -153,6 +155,16 @@ def change_rv_size( Expand the existing size by `new_size`. """ + # Check the dimensionality of the `new_size` kwarg + new_size_ndim = np.ndim(new_size) + if new_size_ndim > 1: + raise ShapeError("The `new_size` must be ≤1-dimensional.", actual=new_size_ndim) + elif new_size_ndim == 0: + new_size = (new_size,) + + # Extract the RV node that is to be resized, together with its inputs, name and tag + 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 @@ -161,10 +173,10 @@ def change_rv_size( if expand: if rv_node.op.ndim_supp == 0 and at.get_vector_length(size) == 0: size = rv_node.op._infer_shape(size, dist_params) - new_size = tuple(at.atleast_1d(new_size)) + tuple(size) + new_size = tuple(new_size) + tuple(size) - # Make sure the new size is a tensor. This helps to not unnecessarily pick - # up a `Cast` in some cases + # Make sure the new size is a tensor. This dtype-aware conversion helps + # to not unnecessarily pick up a `Cast` in some cases (see #4652). new_size = at.as_tensor(new_size, ndim=1, dtype="int64") new_rv_node = rv_node.op.make_node(rng, new_size, dtype, *dist_params) diff --git a/pymc3/tests/test_aesaraf.py b/pymc3/tests/test_aesaraf.py index 90b28ac6c8c..24b59780089 100644 --- a/pymc3/tests/test_aesaraf.py +++ b/pymc3/tests/test_aesaraf.py @@ -41,6 +41,7 @@ take_along_axis, walk_model, ) +from pymc3.exceptions import ShapeError from pymc3.vartypes import int_types FLOATX = str(aesara.config.floatX) @@ -53,6 +54,11 @@ def test_change_rv_size(): assert rv.ndim == 1 assert tuple(rv.shape.eval()) == (2,) + with pytest.raises(ShapeError, match="must be ≤1-dimensional"): + change_rv_size(rv, new_size=[[2, 3]]) + with pytest.raises(ShapeError, match="must be ≤1-dimensional"): + change_rv_size(rv, new_size=at.as_tensor_variable([[2, 3], [4, 5]])) + rv_new = change_rv_size(rv, new_size=(3,), expand=True) assert rv_new.ndim == 2 assert tuple(rv_new.shape.eval()) == (3, 2) diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py index 18a864cb11c..f9844ce5986 100644 --- a/pymc3/tests/test_distributions_random.py +++ b/pymc3/tests/test_distributions_random.py @@ -187,11 +187,10 @@ def get_random_variable(self, shape, with_vector_params=False, name=None): @staticmethod def sample_random_variable(random_variable, size): - """Draws samples from a RandomVariable using its .random() method.""" - if size is None: - return random_variable.eval() - else: - return change_rv_size(random_variable, size, expand=True).eval() + """ Draws samples from a RandomVariable. """ + if size: + random_variable = change_rv_size(random_variable, size, expand=True) + return random_variable.eval() @pytest.mark.parametrize("size", [None, (), 1, (1,), 5, (4, 5)], ids=str) @pytest.mark.parametrize("shape", [None, ()], ids=str) From b9b0b23219d68f494ad7e1de6a5ffe3f57ab0a6c Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Fri, 14 May 2021 21:25:37 +0200 Subject: [PATCH 4/9] Introduce ShapeWarning --- pymc3/exceptions.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pymc3/exceptions.py b/pymc3/exceptions.py index 8f67969709c..cc7f4a5674d 100644 --- a/pymc3/exceptions.py +++ b/pymc3/exceptions.py @@ -17,6 +17,7 @@ "IncorrectArgumentsError", "TraceDirectoryError", "ImputationWarning", + "ShapeWarning", "ShapeError", ] @@ -41,6 +42,12 @@ class ImputationWarning(UserWarning): pass +class ShapeWarning(UserWarning): + """ Something that could lead to shape problems down the line. """ + + pass + + class ShapeError(Exception): """Error that the shape of a variable is incorrect.""" From 124b33c0462bcbd93411c6e576032e72519d2e8f Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Sat, 15 May 2021 16:52:44 +0200 Subject: [PATCH 5/9] Refactor test to use InferenceData --- pymc3/tests/test_data_container.py | 69 ++++++++++++++++++++++++++---- 1 file changed, 61 insertions(+), 8 deletions(-) diff --git a/pymc3/tests/test_data_container.py b/pymc3/tests/test_data_container.py index 99d6b693c7d..93acae3057c 100644 --- a/pymc3/tests/test_data_container.py +++ b/pymc3/tests/test_data_container.py @@ -18,10 +18,13 @@ from aesara import shared from aesara.tensor.sharedvar import ScalarSharedVariable +from aesara.tensor.var import TensorVariable import pymc3 as pm +from pymc3.aesaraf import floatX from pymc3.distributions import logpt +from pymc3.exceptions import ShapeError from pymc3.tests.helpers import SeededTest @@ -159,22 +162,40 @@ def test_shared_data_as_rv_input(self): """ with pm.Model() as m: x = pm.Data("x", [1.0, 2.0, 3.0]) - _ = pm.Normal("y", mu=x, size=3) - trace = pm.sample( - chains=1, return_inferencedata=False, compute_convergence_checks=False + y = pm.Normal("y", mu=x, size=(2, 3)) + assert y.eval().shape == (2, 3) + idata = pm.sample( + chains=1, + tune=500, + draws=550, + return_inferencedata=True, + compute_convergence_checks=False, ) + samples = idata.posterior["y"] + assert samples.shape == (1, 550, 2, 3) np.testing.assert_allclose(np.array([1.0, 2.0, 3.0]), x.get_value(), atol=1e-1) - np.testing.assert_allclose(np.array([1.0, 2.0, 3.0]), trace["y"].mean(0), atol=1e-1) + np.testing.assert_allclose( + np.array([1.0, 2.0, 3.0]), samples.mean(("chain", "draw", "y_dim_0")), atol=1e-1 + ) with m: pm.set_data({"x": np.array([2.0, 4.0, 6.0])}) - trace = pm.sample( - chains=1, return_inferencedata=False, compute_convergence_checks=False + assert y.eval().shape == (2, 3) + idata = pm.sample( + chains=1, + tune=500, + draws=620, + return_inferencedata=True, + compute_convergence_checks=False, ) + samples = idata.posterior["y"] + assert samples.shape == (1, 620, 2, 3) np.testing.assert_allclose(np.array([2.0, 4.0, 6.0]), x.get_value(), atol=1e-1) - np.testing.assert_allclose(np.array([2.0, 4.0, 6.0]), trace["y"].mean(0), atol=1e-1) + np.testing.assert_allclose( + np.array([2.0, 4.0, 6.0]), samples.mean(("chain", "draw", "y_dim_0")), atol=1e-1 + ) def test_shared_scalar_as_rv_input(self): # See https://github.com/pymc-devs/pymc3/issues/3139 @@ -217,7 +238,7 @@ def test_set_data_to_non_data_container_variables(self): ) with pytest.raises(TypeError) as error: pm.set_data({"beta": [1.1, 2.2, 3.3]}, model=model) - error.match("defined as `pymc3.Data` inside the model") + error.match("The variable `beta` must be a `SharedVariable`") @pytest.mark.xfail(reason="Depends on ModelGraph") def test_model_to_graphviz_for_model_with_data_container(self): @@ -283,6 +304,38 @@ def test_explicit_coords(self): assert isinstance(pmodel.dim_lengths["columns"], ScalarSharedVariable) assert pmodel.dim_lengths["columns"].eval() == 7 + def test_symbolic_coords(self): + """ + In v4 dimensions can be created without passing coordinate values. + Their lengths are then automatically linked to the corresponding Tensor dimension. + """ + with pm.Model() as pmodel: + intensity = pm.Data("intensity", np.ones((2, 3)), dims=("row", "column")) + assert "row" in pmodel.dim_lengths + assert "column" in pmodel.dim_lengths + assert isinstance(pmodel.dim_lengths["row"], TensorVariable) + assert isinstance(pmodel.dim_lengths["column"], TensorVariable) + assert pmodel.dim_lengths["row"].eval() == 2 + assert pmodel.dim_lengths["column"].eval() == 3 + + intensity.set_value(floatX(np.ones((4, 5)))) + assert pmodel.dim_lengths["row"].eval() == 4 + assert pmodel.dim_lengths["column"].eval() == 5 + + def test_no_resize_of_implied_dimensions(self): + with pm.Model() as pmodel: + # Imply a dimension through RV params + pm.Normal("n", mu=[1, 2, 3], dims="city") + # _Use_ the dimension for a data variable + inhabitants = pm.Data("inhabitants", [100, 200, 300], dims="city") + + # Attempting to re-size the dimension through the data variable would + # cause shape problems in InferenceData conversion, because the RV remains (3,). + with pytest.raises( + ShapeError, match="was initialized from 'n' which is not a shared variable" + ): + pmodel.set_data("inhabitants", [1, 2, 3, 4]) + def test_implicit_coords_series(self): ser_sales = pd.Series( data=np.random.randint(low=0, high=30, size=22), From c7d477c3d006c0ff8d7ae69282efc5e84909de5a Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Sat, 15 May 2021 16:53:15 +0200 Subject: [PATCH 6/9] Implement backwards-compatble shape and Ellipsis-enabled dims Co-authored-by: Thomas Wiecki Co-authored-by: Ricardo --- RELEASE-NOTES.md | 4 + pymc3/distributions/distribution.py | 365 ++++++++++++++++++++++++---- pymc3/model.py | 54 ++-- pymc3/tests/test_model.py | 37 +-- pymc3/tests/test_ode.py | 5 +- pymc3/tests/test_shape_handling.py | 258 ++++++++++++++++++-- 6 files changed, 604 insertions(+), 119 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 3bdf6272dc7..e08f3010ef3 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -10,6 +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 [#4696](https://github.com/pymc-devs/pymc3/pull/4696)): + - With `shape` the length of dimensions must be given numerically or as scalar Aesara `Variables`. Numeric entries in `shape` restrict 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. An `Ellipsis` (`...`) in the last position of `dims` can be used as short-hand notation for implied dimensions. + - The `size` kwarg behaves like it does in Aesara/NumPy. For univariate RVs it is the same as `shape`, but for multivariate RVs it depends on how the RV implements broadcasting to dimensionality greater than `RVOp.ndim_supp`. - Add `logcdf` method to Kumaraswamy distribution (see [#4706](https://github.com/pymc-devs/pymc3/pull/4706)). - ... diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index 41954e83dd4..562332a59d4 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -19,22 +19,21 @@ import warnings from abc import ABCMeta -from typing import TYPE_CHECKING +from typing import Optional, Sequence, Tuple, Union +import aesara +import aesara.tensor as at import dill +import numpy as np +from aesara.graph.basic import Variable from aesara.tensor.random.op import RandomVariable from aesara.tensor.random.var import RandomStateSharedVariable +from aesara.tensor.var import TensorVariable +from pymc3.aesaraf import change_rv_size, pandas_to_array from pymc3.distributions import _logcdf, _logp - -if TYPE_CHECKING: - from typing import Optional, Callable - -import aesara -import aesara.graph.basic -import aesara.tensor as at - +from pymc3.exceptions import ShapeError, ShapeWarning from pymc3.util import UNSET, get_repr_for_variable from pymc3.vartypes import string_types @@ -52,6 +51,20 @@ PLATFORM = sys.platform +# User-provided can be lazily specified as scalars +Shape = Union[int, TensorVariable, Sequence[Union[int, TensorVariable, type(Ellipsis)]]] +Dims = Union[str, Sequence[Union[str, None, type(Ellipsis)]]] +Size = Union[int, TensorVariable, Sequence[Union[int, TensorVariable]]] + +# After conversion to vectors +WeakShape = Union[TensorVariable, Tuple[Union[int, TensorVariable, type(Ellipsis)], ...]] +WeakDims = Tuple[Union[str, None, type(Ellipsis)], ...] + +# After Ellipsis were substituted +StrongShape = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] +StrongDims = Sequence[Union[str, None]] +StrongSize = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] + class _Unpickling: pass @@ -107,13 +120,187 @@ def logcdf(op, var, rvs_to_values, *dist_params, **kwargs): return new_cls +def _convert_dims(dims: Dims) -> Optional[WeakDims]: + """ Process a user-provided dims variable into None or a valid dims tuple. """ + if dims is None: + return None + + if isinstance(dims, str): + dims = (dims,) + elif isinstance(dims, (list, tuple)): + dims = tuple(dims) + else: + raise ValueError(f"The `dims` parameter must be a tuple, str or list. Actual: {type(dims)}") + + if any(d == Ellipsis for d in dims[:-1]): + raise ValueError(f"Ellipsis in `dims` may only appear in the last position. Actual: {dims}") + + return dims + + +def _convert_shape(shape: Shape) -> Optional[WeakShape]: + """ Process a user-provided shape variable into None or a valid shape object. """ + if shape is None: + return None + + if isinstance(shape, int) or (isinstance(shape, TensorVariable) and shape.ndim == 0): + shape = (shape,) + elif isinstance(shape, (list, tuple)): + shape = tuple(shape) + else: + raise ValueError( + f"The `shape` parameter must be a tuple, TensorVariable, int or list. Actual: {type(shape)}" + ) + + if isinstance(shape, tuple) and any(s == Ellipsis for s in shape[:-1]): + raise ValueError( + f"Ellipsis in `shape` may only appear in the last position. Actual: {shape}" + ) + + return shape + + +def _convert_size(size: Size) -> Optional[StrongSize]: + """ Process a user-provided size variable into None or a valid size object. """ + if size is None: + return None + + if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim == 0): + size = (size,) + elif isinstance(size, (list, tuple)): + size = tuple(size) + else: + raise ValueError( + f"The `size` parameter must be a tuple, TensorVariable, int or list. Actual: {type(size)}" + ) + + if isinstance(size, tuple) and Ellipsis in size: + raise ValueError(f"The `size` parameter cannot contain an Ellipsis. Actual: {size}") + + return size + + +def _resize_from_dims( + dims: WeakDims, ndim_implied: int, model +) -> Tuple[int, StrongSize, StrongDims]: + """Determines a potential resize shape from a `dims` tuple. + + Parameters + ---------- + dims : array-like + A vector of dimension names, None or Ellipsis. + ndim_implied : int + Number of RV dimensions that were implied from its inputs alone. + model : pm.Model + The current model on stack. + + Returns + ------- + ndim_resize : int + Number of dimensions that should be added through resizing. + resize_shape : array-like + The shape of the new dimensions. + """ + if Ellipsis in dims: + # Auto-complete the dims tuple to the full length. + # We don't have a way to know the names of implied + # dimensions, so they will be `None`. + dims = (*dims[:-1], *[None] * ndim_implied) + + ndim_resize = len(dims) - ndim_implied + + # All resize dims must be known already (numerically or symbolically). + unknowndim_resize_dims = set(dims[:ndim_resize]) - set(model.dim_lengths) + if unknowndim_resize_dims: + raise KeyError( + f"Dimensions {unknowndim_resize_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_shape = tuple(model.dim_lengths[dname] for dname in dims[:ndim_resize]) + return ndim_resize, resize_shape, dims + + +def _resize_from_observed( + observed, ndim_implied: int +) -> Tuple[int, StrongSize, Union[np.ndarray, Variable]]: + """Determines a potential resize shape from observations. + + Parameters + ---------- + observed : scalar, array-like + The value of the `observed` kwarg to the RV creation. + ndim_implied : int + Number of RV dimensions that were implied from its inputs alone. + + Returns + ------- + ndim_resize : int + Number of dimensions that should be added through resizing. + resize_shape : array-like + The shape of the new dimensions. + observed : scalar, array-like + Observations as numpy array or `Variable`. + """ + if not hasattr(observed, "shape"): + observed = pandas_to_array(observed) + ndim_resize = observed.ndim - ndim_implied + resize_shape = tuple(observed.shape[d] for d in range(ndim_resize)) + return ndim_resize, resize_shape, observed + + class Distribution(metaclass=DistributionMeta): """Statistical distribution""" rv_class = None rv_op = None - def __new__(cls, name, *args, **kwargs): + def __new__( + cls, + name: str, + *args, + rng=None, + dims: Optional[Dims] = None, + initval=None, + observed=None, + total_size=None, + transform=UNSET, + **kwargs, + ) -> RandomVariable: + """Adds a RandomVariable corresponding to a PyMC3 distribution to the current model. + + Note that all remaining kwargs must be compatible with ``.dist()`` + + Parameters + ---------- + cls : type + A PyMC3 distribution. + name : str + Name for the new model variable. + rng : optional + Random number generator to use with the RandomVariable. + dims : tuple, optional + A tuple of dimension names known to the model. + initval : optional + Test value to be attached to the output RV. + Must match its shape exactly. + observed : optional + Observed data to be passed when registering the random variable in the model. + See ``Model.register_rv``. + total_size : float, optional + See ``Model.register_rv``. + transform : optional + See ``Model.register_rv``. + **kwargs + Keyword arguments that will be forwarded to ``.dist()``. + Most prominently: ``shape`` and ``size`` + + Returns + ------- + rv : RandomVariable + The created RV, registered in the Model. + """ + try: from pymc3.model import Model @@ -126,51 +313,84 @@ def __new__(cls, name, *args, **kwargs): "for a standalone distribution." ) - rng = kwargs.pop("rng", None) - - if rng is None: - rng = model.next_rng() + if "testval" in kwargs: + initval = kwargs.pop("testval") + warnings.warn( + "The `testval` argument is deprecated; use `initval`.", + DeprecationWarning, + stacklevel=2, + ) if not isinstance(name, string_types): raise TypeError(f"Name needs to be a string but got: {name}") - data = kwargs.pop("observed", None) - - total_size = kwargs.pop("total_size", None) - - dims = kwargs.pop("dims", None) - - if "shape" in kwargs: - raise DeprecationWarning("The `shape` keyword is deprecated; use `size`.") - - testval = kwargs.pop("testval", None) + if rng is None: + rng = model.next_rng() - if testval is not None: - warnings.warn( - "The `testval` argument is deprecated; use `initval`.", - DeprecationWarning, - stacklevel=2, + if dims is not None and "shape" in kwargs: + raise ValueError( + f"Passing both `dims` ({dims}) and `shape` ({kwargs['shape']}) is not supported!" + ) + if dims is not None and "size" in kwargs: + raise ValueError( + f"Passing both `dims` ({dims}) and `size` ({kwargs['size']}) is not supported!" ) + dims = _convert_dims(dims) - initval = kwargs.pop("initval", testval) + # 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) + ndim_actual = rv_out.ndim + resize_shape = None - transform = kwargs.pop("transform", UNSET) + # `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: + ndim_resize, resize_shape, dims = _resize_from_dims(dims, ndim_actual, model) + elif observed is not None: + ndim_resize, resize_shape, observed = _resize_from_observed(observed, ndim_actual) - rv_out = cls.dist(*args, rng=rng, **kwargs) + if resize_shape: + # A batch size was specified through `dims`, or implied by `observed`. + rv_out = change_rv_size(rv_var=rv_out, new_size=resize_shape, expand=True) - if testval is not None: - rv_out.tag.test_value = testval + if initval is not None: + # Assigning the testval earlier causes trouble because the RV may not be created with the final shape already. + rv_out.tag.test_value = initval - return model.register_rv( - rv_out, name, data, total_size, dims=dims, transform=transform, initval=initval - ) + return model.register_rv(rv_out, name, observed, total_size, dims=dims, transform=transform) @classmethod - def dist(cls, dist_params, rng=None, **kwargs): + def dist( + cls, + dist_params, + *, + shape: Optional[Shape] = None, + size: Optional[Size] = None, + initval=None, + **kwargs, + ) -> RandomVariable: + """Creates a RandomVariable corresponding to the `cls` distribution. - testval = kwargs.pop("testval", None) + Parameters + ---------- + dist_params : array-like + The inputs to the `RandomVariable` `Op`. + shape : int, tuple, Variable, optional + A tuple of sizes for each dimension of the new RV. + size : int, tuple, Variable, optional + For creating the RV like in Aesara/NumPy. + initival : optional + Test value to be attached to the output RV. + Must match its shape exactly. - if testval is not None: + Returns + ------- + rv : RandomVariable + The created RV. + """ + if "testval" in kwargs: + initval = kwargs.pop("testval") warnings.warn( "The `testval` argument is deprecated. " "Use `initval` to set initial values for a `Model`; " @@ -179,12 +399,65 @@ def dist(cls, dist_params, rng=None, **kwargs): DeprecationWarning, stacklevel=2, ) + if "dims" in kwargs: + raise NotImplementedError("The use of a `.dist(dims=...)` API is not supported.") + if shape is not None and size is not None: + raise ValueError( + f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!" + ) + shape = _convert_shape(shape) + size = _convert_size(size) + + ndim_supp = cls.rv_op.ndim_supp + ndim_expected = None + ndim_batch = None + create_size = None + + if shape is not None: + ndim_expected = len(tuple(shape)) + ndim_batch = ndim_expected - ndim_supp + create_size = tuple(shape)[:ndim_batch] + elif size is not None: + ndim_expected = ndim_supp + len(tuple(size)) + ndim_batch = ndim_expected - ndim_supp + create_size = size + + # Create the RV with a `size` right away. + # This is not necessarily the final result. + rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs) + ndim_actual = rv_out.ndim + ndims_unexpected = ndim_actual != ndim_expected + + if shape is not None and ndims_unexpected: + # This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)). + # Recreate the RV without passing `size` to created it with just the implied dimensions. + rv_out = cls.rv_op(*dist_params, size=None, **kwargs) + + # Now resize by the "extra" dimensions that were not implied from support and parameters + if rv_out.ndim < ndim_expected: + expand_shape = shape[: ndim_expected - rv_out.ndim] + rv_out = change_rv_size(rv_var=rv_out, new_size=expand_shape, expand=True) + if not rv_out.ndim == ndim_expected: + raise ShapeError( + f"Failed to create the RV with the expected dimensionality. " + f"This indicates a severe problem. Please open an issue.", + actual=ndim_actual, + expected=ndim_batch + ndim_supp, + ) - rv_var = cls.rv_op(*dist_params, rng=rng, **kwargs) + # Warn about the edge cases where the RV Op creates more dimensions than + # it should based on `size` and `RVOp.ndim_supp`. + if size is not None and ndims_unexpected: + warnings.warn( + f"You may have expected a ({len(tuple(size))}+{ndim_supp})-dimensional RV, but the resulting RV will be {ndim_actual}-dimensional." + ' To silence this warning use `warnings.simplefilter("ignore", pm.ShapeWarning)`.', + ShapeWarning, + ) + rng = kwargs.pop("rng", None) if ( - rv_var.owner - and isinstance(rv_var.owner.op, RandomVariable) + rv_out.owner + and isinstance(rv_out.owner.op, RandomVariable) and isinstance(rng, RandomStateSharedVariable) and not getattr(rng, "default_update", None) ): @@ -194,11 +467,11 @@ def dist(cls, dist_params, rng=None, **kwargs): # Without it, the `RandomVariable`s could not be optimized to allow # in-place RNG updates, forcing all sample results from compiled # functions to be the same on repeated evaluations. - new_rng = rv_var.owner.outputs[0] - rv_var.update = (rng, new_rng) + new_rng = rv_out.owner.outputs[0] + rv_out.update = (rng, new_rng) rng.default_update = new_rng - return rv_var + return rv_out def _distr_parameters_for_repr(self): """Return the names of the parameters for this distribution (e.g. "mu" diff --git a/pymc3/model.py b/pymc3/model.py index cc80df504ab..d27763a4f5c 100644 --- a/pymc3/model.py +++ b/pymc3/model.py @@ -48,7 +48,6 @@ from pandas import Series from pymc3.aesaraf import ( - change_rv_size, compile_rv_inplace, gradient, hessian, @@ -1061,14 +1060,18 @@ def set_data( New values for the shared variable. coords : optional, dict New coordinate values for dimensions of the shared variable. - Must be provided for all named dimensions that change in length. + Must be provided for all named dimensions that change in length + and already have coordinate values. """ shared_object = self[name] if not isinstance(shared_object, SharedVariable): raise TypeError( - f"The variable `{name}` must be defined as `pymc3.Data` inside the model to allow updating. " + f"The variable `{name}` must be a `SharedVariable` (e.g. `pymc3.Data`) to allow updating. " f"The current type is: {type(shared_object)}" ) + + if isinstance(values, list): + values = np.array(values) values = pandas_to_array(values) dims = self.RV_dims.get(name, None) or () coords = coords or {} @@ -1090,10 +1093,11 @@ def set_data( # Reject resizing if we already know that it would create shape problems. # NOTE: If there are multiple pm.Data containers sharing this dim, but the user only # changes the values for one of them, they will run into shape problems nonetheless. - if not isinstance(length_tensor, ScalarSharedVariable) and length_changed: + length_belongs_to = length_tensor.owner.inputs[0].owner.inputs[0] + if not isinstance(length_belongs_to, SharedVariable) and length_changed: raise ShapeError( - f"Resizing dimension {dname} with values of length {new_length} would lead to incompatibilities, " - f"because the dimension was not initialized from a shared variable. " + f"Resizing dimension '{dname}' with values of length {new_length} would lead to incompatibilities, " + f"because the dimension was initialized from '{length_belongs_to}' which is not a shared variable. " f"Check if the dimension was defined implicitly before the shared variable '{name}' was created, " f"for example by a model variable.", actual=new_length, @@ -1130,6 +1134,7 @@ def register_rv( ---------- rv_var: TensorVariable name: str + Intended name for the model variable. data: array_like (optional) If data is provided, the variable is observed. If None, the variable is unobserved. @@ -1150,6 +1155,13 @@ def register_rv( rv_var.name = name rv_var.tag.total_size = total_size + # Associate previously unknown dimension names with + # the length of the corresponding RV dimension. + if dims is not None: + for d, dname in enumerate(dims): + if not dname in self.dim_lengths: + self.add_coord(dname, values=None, length=rv_var.shape[d]) + if data is None: self.free_RVs.append(rv_var) self.create_value_var(rv_var, transform) @@ -1161,11 +1173,13 @@ def register_rv( and not isinstance(data, (GenTensorVariable, Minibatch)) and data.owner is not None ): - raise TypeError("Observed data cannot consist of symbolic variables.") + raise TypeError( + "Variables that depend on other nodes cannot be used for observed data." + f"The data variable was: {data}" + ) - # `rv_var` is potentially a new variable (e.g. the original - # variable could have its size changed to match the data, or be a - # new graph that accounts for missing data) + # `rv_var` is potentially changed by `make_obs_var`, + # for example into a new graph for imputation of missing data. rv_var = self.make_obs_var(rv_var, data, dims, transform) return rv_var @@ -1179,6 +1193,7 @@ def make_obs_var( ========== rv_var The random variable that is observed. + Its dimensionality must be compatible with the data already. data The observed data. dims: tuple @@ -1190,21 +1205,10 @@ def make_obs_var( name = rv_var.name data = pandas_to_array(data).astype(rv_var.dtype) - # The shapes of the observed random variable and its data might not - # match. We need need to update the observed random variable's `size` - # (i.e. number of samples) so that it matches the data. - - # Setting `size` produces a random variable with shape `size + - # support_shape`, where `len(support_shape) == op.ndim_supp`, we need - # to disregard the last `op.ndim_supp`-many dimensions when we - # determine the appropriate `size` value from `data.shape`. - ndim_supp = rv_var.owner.op.ndim_supp - if ndim_supp > 0: - new_size = data.shape[:-ndim_supp] - else: - new_size = data.shape - - rv_var = change_rv_size(rv_var, new_size) + if data.ndim != rv_var.ndim: + raise ShapeError( + "Dimensionality of data and RV don't match.", actual=data.ndim, expected=rv_var.ndim + ) if aesara.config.compute_test_value != "off": test_value = getattr(rv_var.tag, "test_value", None) diff --git a/pymc3/tests/test_model.py b/pymc3/tests/test_model.py index 53ab66af217..d51952b58ac 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 @@ -461,33 +462,33 @@ def test_make_obs_var(): # Create a fake model and fake distribution to be used for the test fake_model = pm.Model() with fake_model: - fake_distribution = pm.Normal.dist(mu=0, sigma=1) - # Create the initval attribute simply for the sake of model testing + fake_distribution = pm.Normal.dist(mu=0, sigma=1, size=(3, 3)) + # 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 + dense_output = fake_model.make_obs_var(fake_distribution, dense_input, None, None) + assert dense_output == fake_distribution + assert isinstance(dense_output.tag.observations, TensorConstant) del fake_model.named_vars[fake_distribution.name] + sparse_output = fake_model.make_obs_var(fake_distribution, sparse_input, None, None) + assert sparse_output == fake_distribution + assert sparse.basic._is_sparse_variable(sparse_output.tag.observations) del fake_model.named_vars[fake_distribution.name] + + # Here the RandomVariable is split into observed/imputed and a Deterministic is returned masked_output = fake_model.make_obs_var(fake_distribution, masked_array_input, None, None) + assert masked_output != fake_distribution assert not isinstance(masked_output, RandomVariable) - - # Ensure that the missing values are appropriately set to None - for func_output in [dense_output, sparse_output]: - assert isinstance(func_output.owner.op, RandomVariable) - - # Ensure that the Aesara variable names are correctly set. - # Note that the output for masked inputs do not have their names set - # to the passed value. - for func_output in [dense_output, sparse_output]: - assert func_output.name == input_name - - # Ensure the that returned functions are all of the correct type - assert isinstance(dense_output.tag.observations, TensorConstant) - assert sparse.basic._is_sparse_variable(sparse_output.tag.observations) - - # Masked output is something weird. Just ensure it has missing values + # Ensure it has missing values assert {"testing_inputs_missing"} == {v.name for v in fake_model.vars} assert {"testing_inputs", "testing_inputs_observed"} == { v.name for v in fake_model.observed_RVs diff --git a/pymc3/tests/test_ode.py b/pymc3/tests/test_ode.py index 70b935fea42..d910e7b6b49 100644 --- a/pymc3/tests/test_ode.py +++ b/pymc3/tests/test_ode.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import sys import aesara import numpy as np @@ -169,9 +168,7 @@ def ode_func_5(y, t, p): np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._sens_ic) -@pytest.mark.xfail( - condition=sys.platform == "win32", reason="https://github.com/pymc-devs/aesara/issues/390" -) +@pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") def test_logp_scalar_ode(): """Test the computation of the log probability for these models""" diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index 141c28dad3f..5bb7aa0288c 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -20,6 +20,11 @@ import pymc3 as pm +from pymc3.distributions.distribution import ( + _convert_dims, + _convert_shape, + _convert_size, +) from pymc3.distributions.shape_utils import ( broadcast_dist_samples_shape, broadcast_dist_samples_to, @@ -28,6 +33,7 @@ shapes_broadcasting, to_tuple, ) +from pymc3.exceptions import ShapeWarning test_shapes = [ (tuple(), (1,), (4,), (5, 4)), @@ -222,30 +228,230 @@ def test_sample_generate_values(fixture_model, fixture_sizes): assert prior[rv.name].shape == size + tuple(rv.distribution.shape) -@pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") -def test_size32_doesnt_break_broadcasting(): - size32 = at.constant([1, 10], dtype="int32") - rv = pm.Normal.dist(0, 1, size=size32) - assert rv.broadcastable == (True, False) - - -def test_observed_with_column_vector(): - """This test is related to https://github.com/pymc-devs/aesara/issues/390 which breaks - broadcastability of column-vector RVs. This unexpected change in type can lead to - incompatibilities during graph rewriting for model.logp evaluation. - """ - with pm.Model() as model: - # The `observed` is a broadcastable column vector - obs = at.as_tensor_variable(np.ones((3, 1), dtype=aesara.config.floatX)) - assert obs.broadcastable == (False, True) - - # Both shapes describe broadcastable volumn vectors - size64 = at.constant([3, 1], dtype="int64") - # But the second shape is upcasted from an int32 vector - cast64 = at.cast(at.constant([3, 1], dtype="int32"), dtype="int64") - - pm.Normal("x_size64", mu=0, sd=1, size=size64, observed=obs) - model.logp() +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 + ): + coords = {} + param_dims = [] + batch_dims = [] + + # Create coordinates corresponding to the parameter shape + for d in param_shape: + dname = f"param_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 + 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) + + with pm.Model(coords=coords) as pmodel: + mu = aesara.shared(np.random.normal(size=param_shape)) + + with pytest.warns(None): + if parametrization == "implicit": + rv = pm.Normal("rv", mu=mu).shape == param_shape + else: + expected_shape = batch_shape + param_shape + if parametrization == "shape": + rv = pm.Normal("rv", mu=mu, shape=batch_shape + param_shape) + assert rv.eval().shape == expected_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 == expected_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 + param_shape) + assert rv.eval().shape == expected_shape + else: + raise NotImplementedError("Invalid test case parametrization.") + + def test_define_dims_on_the_fly(self): + with pm.Model() as pmodel: + agedata = aesara.shared(np.array([10, 20, 30])) + + # Associate the "patient" dim with an implied dimension + age = pm.Normal("age", agedata, dims=("patient",)) + assert "patient" in pmodel.dim_lengths + assert pmodel.dim_lengths["patient"].eval() == 3 + + # Use the dim to replicate a new RV + effect = pm.Normal("effect", 0, dims=("patient",)) + assert effect.ndim == 1 + assert effect.eval().shape == (3,) + + # Now change the length of the implied dimension + agedata.set_value([1, 2, 3, 4]) + # The change should propagate all the way through + assert effect.eval().shape == (4,) + + @pytest.mark.xfail(reason="Simultaneous use of size and dims is not implemented") + def test_data_defined_size_dimension_can_register_dimname(self): + with pm.Model() as pmodel: + x = pm.Data("x", [[1, 2, 3, 4]], dims=("first", "second")) + 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")) + assert "third" in pmodel.dim_lengths + assert y.eval().shape() == (2, 1, 4) + + def test_can_resize_data_defined_size(self): + with pm.Model() as pmodel: + x = pm.Data("x", [[1, 2, 3, 4]], dims=("first", "second")) + y = pm.Normal("y", mu=0, dims=("first", "second")) + z = pm.Normal("z", mu=y, observed=np.ones((1, 4))) + assert x.eval().shape == (1, 4) + assert y.eval().shape == (1, 4) + assert z.eval().shape == (1, 4) + assert "first" in pmodel.dim_lengths + assert "second" in pmodel.dim_lengths + pmodel.set_data("x", [[1, 2], [3, 4], [5, 6]]) + assert x.eval().shape == (3, 2) + assert y.eval().shape == (3, 2) + assert z.eval().shape == (3, 2) + + @pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") + def test_size32_doesnt_break_broadcasting(): + size32 = at.constant([1, 10], dtype="int32") + rv = pm.Normal.dist(0, 1, size=size32) + assert rv.broadcastable == (True, False) + + @pytest.mark.xfail(reason="https://github.com/pymc-devs/aesara/issues/390") + def test_observed_with_column_vector(self): + """This test is related to https://github.com/pymc-devs/aesara/issues/390 which breaks + broadcastability of column-vector RVs. This unexpected change in type can lead to + incompatibilities during graph rewriting for model.logp evaluation. + """ + with pm.Model() as model: + # The `observed` is a broadcastable column vector + obs = at.as_tensor_variable(np.ones((3, 1), dtype=aesara.config.floatX)) + assert obs.broadcastable == (False, True) + + # Both shapes describe broadcastable volumn vectors + size64 = at.constant([3, 1], dtype="int64") + # But the second shape is upcasted from an int32 vector + cast64 = at.cast(at.constant([3, 1], dtype="int32"), dtype="int64") + + pm.Normal("size64", mu=0, sd=1, size=size64, observed=obs) + pm.Normal("shape64", mu=0, sd=1, shape=size64, observed=obs) + model.logp() + + pm.Normal("size_cast64", mu=0, sd=1, size=cast64, observed=obs) + pm.Normal("shape_cast64", mu=0, sd=1, shape=cast64, observed=obs) + model.logp() + + def test_dist_api_works(self): + mu = aesara.shared(np.array([1, 2, 3])) + with pytest.raises(NotImplementedError, match="API is not 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=(3,)).eval().shape == (3,) + assert pm.Normal.dist(mu=mu, size=(4, 3)).eval().shape == (4, 3) + + def test_mvnormal_shape_size_difference(self): + # Parameters add one batch dimension (4), shape is what you'd expect. + # Under the hood the shape(4, 3) becomes size=(4,) and the RV is initially + # created as (4, 4, 3). The internal ndim-check then recreates it with size=None. + rv = pm.MvNormal.dist(mu=np.ones((4, 3)), cov=np.eye(3), shape=(4, 3)) + assert rv.ndim == 2 + assert tuple(rv.shape.eval()) == (4, 3) + + # shape adds two dimensions (5, 4) + # Under the hood the shape=(5, 4, 3) becomes size=(5, 4). + # The RV is created as (5, 4, 3) right away. + rv = pm.MvNormal.dist(mu=[1, 2, 3], cov=np.eye(3), shape=(5, 4, 3)) + assert rv.ndim == 3 + assert tuple(rv.shape.eval()) == (5, 4, 3) + + # parameters add 1 batch dimension (4), shape adds another (5) + # Under the hood the shape=(5, 4, 3) becomes size=(5, 4) + # The RV is initially created as (5, 4, 3, 4, 3) and then recreated and resized. + rv = pm.MvNormal.dist(mu=np.ones((4, 3)), cov=np.eye(3), shape=(5, 4, 3)) + assert rv.ndim == 3 + assert tuple(rv.shape.eval()) == (5, 4, 3) + + # rv = pm.MvNormal.dist(mu=np.ones((4, 3, 2)), cov=np.eye(2), shape=(6, 5, ...)) + # assert rv.ndim == 5 + # assert tuple(rv.shape.eval()) == (6, 5, 4, 3, 2) + + with pytest.warns(None): + rv = pm.MvNormal.dist(mu=[1, 2, 3], cov=np.eye(3), size=(5, 4)) + assert tuple(rv.shape.eval()) == (5, 4, 3) + + # When using `size` the API behaves like Aesara/NumPy + with pytest.warns( + ShapeWarning, + match=r"You may have expected a \(2\+1\)-dimensional RV, but the resulting RV will be 5-dimensional", + ): + rv = pm.MvNormal.dist(mu=np.ones((5, 4, 3)), cov=np.eye(3), size=(5, 4)) + assert tuple(rv.shape.eval()) == (5, 4, 5, 4, 3) + + def test_convert_dims(self): + assert _convert_dims(dims="town") == ("town",) + with pytest.raises(ValueError, match="must be a tuple, str or list"): + _convert_dims(3) + with pytest.raises(ValueError, match="may only appear in the last position"): + _convert_dims(dims=(..., "town")) + + def test_convert_shape(self): + assert _convert_shape(5) == (5,) + with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): + _convert_shape(shape="notashape") + with pytest.raises(ValueError, match="may only appear in the last position"): + _convert_shape(shape=(3, ..., 2)) + + def test_convert_size(self): + assert _convert_size(7) == (7,) + with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): + _convert_size(size="notasize") + with pytest.raises(ValueError, match="cannot contain"): + _convert_size(size=(3, ...)) + + def test_lazy_flavors(self): + assert pm.Uniform.dist(2, [4, 5], size=[3, 2]).eval().shape == (3, 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("n1", mu=[1, 2], dims="town").eval().shape == (2,) + assert pm.Normal("n2", mu=[1, 2], dims=["town"]).eval().shape == (2,) + + def test_invalid_flavors(self): + with pytest.raises(ValueError, match="Passing both"): + pm.Normal.dist(0, 1, shape=(3,), size=(3,)) - pm.Normal("x_cast64", mu=0, sd=1, size=cast64, observed=obs) - model.logp() + with pm.Model(): + with pytest.raises(ValueError, match="Passing both"): + pm.Normal("n", shape=(2,), dims=("town",)) + with pytest.raises(ValueError, match="Passing both"): + pm.Normal("n", dims=("town",), size=(2,)) From 46c225e4bc440a2e22b9b6aaaf5bf8f1c78ffa9e Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Sun, 16 May 2021 14:20:19 +0200 Subject: [PATCH 7/9] Add Ellipsis-support for the shape kwarg --- RELEASE-NOTES.md | 3 +- pymc3/distributions/distribution.py | 50 ++++++++++++++++++----------- pymc3/tests/test_shape_handling.py | 16 ++++----- 3 files changed, 42 insertions(+), 27 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index e08f3010ef3..f1daf593aa3 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -12,8 +12,9 @@ - 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 [#4696](https://github.com/pymc-devs/pymc3/pull/4696)): - With `shape` the length of dimensions must be given numerically or as scalar Aesara `Variables`. Numeric entries in `shape` restrict 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. An `Ellipsis` (`...`) in the last position of `dims` can be used as short-hand notation for implied dimensions. + - `dims` keeps model variables re-sizeable (for example through `pm.Data`) and leads to well defined coordinates in `InferenceData` objects. - The `size` kwarg behaves like it does in Aesara/NumPy. For univariate RVs it is the same as `shape`, but for multivariate RVs it depends on how the RV implements broadcasting to dimensionality greater than `RVOp.ndim_supp`. + - An `Ellipsis` (`...`) in the last position of `shape` or `dims` can be used as short-hand notation for implied dimensions. - Add `logcdf` method to Kumaraswamy distribution (see [#4706](https://github.com/pymc-devs/pymc3/pull/4706)). - ... diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index 562332a59d4..a334a8627eb 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -378,6 +378,9 @@ def dist( The inputs to the `RandomVariable` `Op`. shape : int, tuple, Variable, optional A tuple of sizes for each dimension of the new RV. + + An Ellipsis (...) may be inserted in the last position to short-hand refer to + all the dimensions that the RV would get if no shape/size/dims were passed at all. size : int, tuple, Variable, optional For creating the RV like in Aesara/NumPy. initival : optional @@ -414,9 +417,16 @@ def dist( create_size = None if shape is not None: - ndim_expected = len(tuple(shape)) - ndim_batch = ndim_expected - ndim_supp - create_size = tuple(shape)[:ndim_batch] + if Ellipsis in shape: + # Ellipsis short-hands all implied dimensions. Therefore + # we don't know how many dimensions to expect. + ndim_expected = ndim_batch = None + # Create the RV with its implied shape and resize later + create_size = None + else: + ndim_expected = len(tuple(shape)) + ndim_batch = ndim_expected - ndim_supp + create_size = tuple(shape)[:ndim_batch] elif size is not None: ndim_expected = ndim_supp + len(tuple(size)) ndim_batch = ndim_expected - ndim_supp @@ -429,21 +439,25 @@ def dist( ndims_unexpected = ndim_actual != ndim_expected if shape is not None and ndims_unexpected: - # This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)). - # Recreate the RV without passing `size` to created it with just the implied dimensions. - rv_out = cls.rv_op(*dist_params, size=None, **kwargs) - - # Now resize by the "extra" dimensions that were not implied from support and parameters - if rv_out.ndim < ndim_expected: - expand_shape = shape[: ndim_expected - rv_out.ndim] - rv_out = change_rv_size(rv_var=rv_out, new_size=expand_shape, expand=True) - if not rv_out.ndim == ndim_expected: - raise ShapeError( - f"Failed to create the RV with the expected dimensionality. " - f"This indicates a severe problem. Please open an issue.", - actual=ndim_actual, - expected=ndim_batch + ndim_supp, - ) + if Ellipsis in shape: + # Resize and we're done! + rv_out = change_rv_size(rv_var=rv_out, new_size=shape[:-1], expand=True) + else: + # This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)). + # Recreate the RV without passing `size` to created it with just the implied dimensions. + rv_out = cls.rv_op(*dist_params, size=None, **kwargs) + + # Now resize by any remaining "extra" dimensions that were not implied from support and parameters + if rv_out.ndim < ndim_expected: + expand_shape = shape[: ndim_expected - rv_out.ndim] + rv_out = change_rv_size(rv_var=rv_out, new_size=expand_shape, expand=True) + if not rv_out.ndim == ndim_expected: + raise ShapeError( + f"Failed to create the RV with the expected dimensionality. " + f"This indicates a severe problem. Please open an issue.", + actual=ndim_actual, + expected=ndim_batch + ndim_supp, + ) # Warn about the edge cases where the RV Op creates more dimensions than # it should based on `size` and `RVOp.ndim_supp`. diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index 5bb7aa0288c..2b75ba765ef 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -236,7 +236,7 @@ class TestShapeDimsSize: [ "implicit", "shape", - # "shape...", + "shape...", "dims", "dims...", "size", @@ -273,9 +273,9 @@ def test_param_and_batch_shape_combos( if parametrization == "shape": rv = pm.Normal("rv", mu=mu, shape=batch_shape + param_shape) assert rv.eval().shape == expected_shape - # elif parametrization == "shape...": - # rv = pm.Normal("rv", mu=mu, shape=(*batch_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 == expected_shape @@ -376,7 +376,7 @@ def test_dist_api_works(self): 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, shape=(7, ...)).eval().shape == (7, 3) assert pm.Normal.dist(mu=mu, size=(3,)).eval().shape == (3,) assert pm.Normal.dist(mu=mu, size=(4, 3)).eval().shape == (4, 3) @@ -402,9 +402,9 @@ def test_mvnormal_shape_size_difference(self): assert rv.ndim == 3 assert tuple(rv.shape.eval()) == (5, 4, 3) - # rv = pm.MvNormal.dist(mu=np.ones((4, 3, 2)), cov=np.eye(2), shape=(6, 5, ...)) - # assert rv.ndim == 5 - # assert tuple(rv.shape.eval()) == (6, 5, 4, 3, 2) + rv = pm.MvNormal.dist(mu=np.ones((4, 3, 2)), cov=np.eye(2), shape=(6, 5, ...)) + assert rv.ndim == 5 + assert tuple(rv.shape.eval()) == (6, 5, 4, 3, 2) with pytest.warns(None): rv = pm.MvNormal.dist(mu=[1, 2, 3], cov=np.eye(3), size=(5, 4)) From 8bdc225ba195b80a8261ca8d54744671474a724c Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Tue, 1 Jun 2021 23:03:50 +0200 Subject: [PATCH 8/9] Separate shape logic into a separate file (#4708) Co-authored-by: Michael Osthege --- pymc3/distributions/distribution.py | 241 ++++-------------------- pymc3/distributions/shape_utils.py | 281 +++++++++++++++++++++++++++- pymc3/tests/test_shape_handling.py | 28 ++- 3 files changed, 326 insertions(+), 224 deletions(-) diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index a334a8627eb..a9f090eabed 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -19,21 +19,29 @@ import warnings from abc import ABCMeta -from typing import Optional, Sequence, Tuple, Union +from typing import Optional import aesara import aesara.tensor as at import dill -import numpy as np -from aesara.graph.basic import Variable from aesara.tensor.random.op import RandomVariable from aesara.tensor.random.var import RandomStateSharedVariable -from aesara.tensor.var import TensorVariable -from pymc3.aesaraf import change_rv_size, pandas_to_array +from pymc3.aesaraf import change_rv_size from pymc3.distributions import _logcdf, _logp -from pymc3.exceptions import ShapeError, ShapeWarning +from pymc3.distributions.shape_utils import ( + Dims, + Shape, + Size, + convert_dims, + convert_shape, + convert_size, + find_size, + maybe_resize, + resize_from_dims, + resize_from_observed, +) from pymc3.util import UNSET, get_repr_for_variable from pymc3.vartypes import string_types @@ -51,20 +59,6 @@ PLATFORM = sys.platform -# User-provided can be lazily specified as scalars -Shape = Union[int, TensorVariable, Sequence[Union[int, TensorVariable, type(Ellipsis)]]] -Dims = Union[str, Sequence[Union[str, None, type(Ellipsis)]]] -Size = Union[int, TensorVariable, Sequence[Union[int, TensorVariable]]] - -# After conversion to vectors -WeakShape = Union[TensorVariable, Tuple[Union[int, TensorVariable, type(Ellipsis)], ...]] -WeakDims = Tuple[Union[str, None, type(Ellipsis)], ...] - -# After Ellipsis were substituted -StrongShape = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] -StrongDims = Sequence[Union[str, None]] -StrongSize = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] - class _Unpickling: pass @@ -120,135 +114,6 @@ def logcdf(op, var, rvs_to_values, *dist_params, **kwargs): return new_cls -def _convert_dims(dims: Dims) -> Optional[WeakDims]: - """ Process a user-provided dims variable into None or a valid dims tuple. """ - if dims is None: - return None - - if isinstance(dims, str): - dims = (dims,) - elif isinstance(dims, (list, tuple)): - dims = tuple(dims) - else: - raise ValueError(f"The `dims` parameter must be a tuple, str or list. Actual: {type(dims)}") - - if any(d == Ellipsis for d in dims[:-1]): - raise ValueError(f"Ellipsis in `dims` may only appear in the last position. Actual: {dims}") - - return dims - - -def _convert_shape(shape: Shape) -> Optional[WeakShape]: - """ Process a user-provided shape variable into None or a valid shape object. """ - if shape is None: - return None - - if isinstance(shape, int) or (isinstance(shape, TensorVariable) and shape.ndim == 0): - shape = (shape,) - elif isinstance(shape, (list, tuple)): - shape = tuple(shape) - else: - raise ValueError( - f"The `shape` parameter must be a tuple, TensorVariable, int or list. Actual: {type(shape)}" - ) - - if isinstance(shape, tuple) and any(s == Ellipsis for s in shape[:-1]): - raise ValueError( - f"Ellipsis in `shape` may only appear in the last position. Actual: {shape}" - ) - - return shape - - -def _convert_size(size: Size) -> Optional[StrongSize]: - """ Process a user-provided size variable into None or a valid size object. """ - if size is None: - return None - - if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim == 0): - size = (size,) - elif isinstance(size, (list, tuple)): - size = tuple(size) - else: - raise ValueError( - f"The `size` parameter must be a tuple, TensorVariable, int or list. Actual: {type(size)}" - ) - - if isinstance(size, tuple) and Ellipsis in size: - raise ValueError(f"The `size` parameter cannot contain an Ellipsis. Actual: {size}") - - return size - - -def _resize_from_dims( - dims: WeakDims, ndim_implied: int, model -) -> Tuple[int, StrongSize, StrongDims]: - """Determines a potential resize shape from a `dims` tuple. - - Parameters - ---------- - dims : array-like - A vector of dimension names, None or Ellipsis. - ndim_implied : int - Number of RV dimensions that were implied from its inputs alone. - model : pm.Model - The current model on stack. - - Returns - ------- - ndim_resize : int - Number of dimensions that should be added through resizing. - resize_shape : array-like - The shape of the new dimensions. - """ - if Ellipsis in dims: - # Auto-complete the dims tuple to the full length. - # We don't have a way to know the names of implied - # dimensions, so they will be `None`. - dims = (*dims[:-1], *[None] * ndim_implied) - - ndim_resize = len(dims) - ndim_implied - - # All resize dims must be known already (numerically or symbolically). - unknowndim_resize_dims = set(dims[:ndim_resize]) - set(model.dim_lengths) - if unknowndim_resize_dims: - raise KeyError( - f"Dimensions {unknowndim_resize_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_shape = tuple(model.dim_lengths[dname] for dname in dims[:ndim_resize]) - return ndim_resize, resize_shape, dims - - -def _resize_from_observed( - observed, ndim_implied: int -) -> Tuple[int, StrongSize, Union[np.ndarray, Variable]]: - """Determines a potential resize shape from observations. - - Parameters - ---------- - observed : scalar, array-like - The value of the `observed` kwarg to the RV creation. - ndim_implied : int - Number of RV dimensions that were implied from its inputs alone. - - Returns - ------- - ndim_resize : int - Number of dimensions that should be added through resizing. - resize_shape : array-like - The shape of the new dimensions. - observed : scalar, array-like - Observations as numpy array or `Variable`. - """ - if not hasattr(observed, "shape"): - observed = pandas_to_array(observed) - ndim_resize = observed.ndim - ndim_implied - resize_shape = tuple(observed.shape[d] for d in range(ndim_resize)) - return ndim_resize, resize_shape, observed - - class Distribution(metaclass=DistributionMeta): """Statistical distribution""" @@ -335,7 +200,7 @@ def __new__( raise ValueError( f"Passing both `dims` ({dims}) and `size` ({kwargs['size']}) is not supported!" ) - dims = _convert_dims(dims) + dims = convert_dims(dims) # 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). @@ -346,9 +211,9 @@ def __new__( # `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: - ndim_resize, resize_shape, dims = _resize_from_dims(dims, ndim_actual, model) + ndim_resize, resize_shape, dims = resize_from_dims(dims, ndim_actual, model) elif observed is not None: - ndim_resize, resize_shape, observed = _resize_from_observed(observed, ndim_actual) + ndim_resize, resize_shape, observed = resize_from_observed(observed, ndim_actual) if resize_shape: # A batch size was specified through `dims`, or implied by `observed`. @@ -408,65 +273,27 @@ def dist( raise ValueError( f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!" ) - shape = _convert_shape(shape) - size = _convert_size(size) - - ndim_supp = cls.rv_op.ndim_supp - ndim_expected = None - ndim_batch = None - create_size = None - - if shape is not None: - if Ellipsis in shape: - # Ellipsis short-hands all implied dimensions. Therefore - # we don't know how many dimensions to expect. - ndim_expected = ndim_batch = None - # Create the RV with its implied shape and resize later - create_size = None - else: - ndim_expected = len(tuple(shape)) - ndim_batch = ndim_expected - ndim_supp - create_size = tuple(shape)[:ndim_batch] - elif size is not None: - ndim_expected = ndim_supp + len(tuple(size)) - ndim_batch = ndim_expected - ndim_supp - create_size = size + shape = convert_shape(shape) + size = convert_size(size) + + create_size, ndim_expected, ndim_batch, ndim_supp = find_size( + shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp + ) # Create the RV with a `size` right away. # This is not necessarily the final result. rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs) - ndim_actual = rv_out.ndim - ndims_unexpected = ndim_actual != ndim_expected - - if shape is not None and ndims_unexpected: - if Ellipsis in shape: - # Resize and we're done! - rv_out = change_rv_size(rv_var=rv_out, new_size=shape[:-1], expand=True) - else: - # This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)). - # Recreate the RV without passing `size` to created it with just the implied dimensions. - rv_out = cls.rv_op(*dist_params, size=None, **kwargs) - - # Now resize by any remaining "extra" dimensions that were not implied from support and parameters - if rv_out.ndim < ndim_expected: - expand_shape = shape[: ndim_expected - rv_out.ndim] - rv_out = change_rv_size(rv_var=rv_out, new_size=expand_shape, expand=True) - if not rv_out.ndim == ndim_expected: - raise ShapeError( - f"Failed to create the RV with the expected dimensionality. " - f"This indicates a severe problem. Please open an issue.", - actual=ndim_actual, - expected=ndim_batch + ndim_supp, - ) - - # Warn about the edge cases where the RV Op creates more dimensions than - # it should based on `size` and `RVOp.ndim_supp`. - if size is not None and ndims_unexpected: - warnings.warn( - f"You may have expected a ({len(tuple(size))}+{ndim_supp})-dimensional RV, but the resulting RV will be {ndim_actual}-dimensional." - ' To silence this warning use `warnings.simplefilter("ignore", pm.ShapeWarning)`.', - ShapeWarning, - ) + rv_out = maybe_resize( + rv_out, + cls.rv_op, + dist_params, + ndim_expected, + ndim_batch, + ndim_supp, + shape, + size, + **kwargs, + ) rng = kwargs.pop("rng", None) if ( diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py index a339ae7c08c..c55fdb30b00 100644 --- a/pymc3/distributions/shape_utils.py +++ b/pymc3/distributions/shape_utils.py @@ -1,4 +1,4 @@ -# Copyright 2020 The PyMC Developers +# Copyright 2021 The PyMC Developers # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,11 +14,22 @@ # -*- coding: utf-8 -*- """ -A collection of common numpy array shape operations needed for broadcasting +A collection of common shape operations needed for broadcasting samples from probability distributions for stochastic nodes in PyMC. """ + +import warnings + +from typing import Optional, Sequence, Tuple, Union + import numpy as np +from aesara.graph.basic import Variable +from aesara.tensor.var import TensorVariable + +from pymc3.aesaraf import change_rv_size, pandas_to_array +from pymc3.exceptions import ShapeError, ShapeWarning + __all__ = [ "to_tuple", "shapes_broadcasting", @@ -397,3 +408,269 @@ def broadcast_dist_samples_to(to_shape, samples, size=None): samples, size=size, must_bcast_with=to_shape, return_out_shape=True ) return [np.broadcast_to(o, to_shape) for o in samples] + + +# User-provided can be lazily specified as scalars +Shape = Union[int, TensorVariable, Sequence[Union[int, TensorVariable, type(Ellipsis)]]] +Dims = Union[str, Sequence[Union[str, None, type(Ellipsis)]]] +Size = Union[int, TensorVariable, Sequence[Union[int, TensorVariable]]] + +# After conversion to vectors +WeakShape = Union[TensorVariable, Tuple[Union[int, TensorVariable, type(Ellipsis)], ...]] +WeakDims = Tuple[Union[str, None, type(Ellipsis)], ...] + +# After Ellipsis were substituted +StrongShape = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] +StrongDims = Sequence[Union[str, None]] +StrongSize = Union[TensorVariable, Tuple[Union[int, TensorVariable], ...]] + + +def convert_dims(dims: Dims) -> Optional[WeakDims]: + """ Process a user-provided dims variable into None or a valid dims tuple. """ + if dims is None: + return None + + if isinstance(dims, str): + dims = (dims,) + elif isinstance(dims, (list, tuple)): + dims = tuple(dims) + else: + raise ValueError(f"The `dims` parameter must be a tuple, str or list. Actual: {type(dims)}") + + if any(d == Ellipsis for d in dims[:-1]): + raise ValueError(f"Ellipsis in `dims` may only appear in the last position. Actual: {dims}") + + return dims + + +def convert_shape(shape: Shape) -> Optional[WeakShape]: + """ Process a user-provided shape variable into None or a valid shape object. """ + if shape is None: + return None + + if isinstance(shape, int) or (isinstance(shape, TensorVariable) and shape.ndim == 0): + shape = (shape,) + elif isinstance(shape, (list, tuple)): + shape = tuple(shape) + else: + raise ValueError( + f"The `shape` parameter must be a tuple, TensorVariable, int or list. Actual: {type(shape)}" + ) + + if isinstance(shape, tuple) and any(s == Ellipsis for s in shape[:-1]): + raise ValueError( + f"Ellipsis in `shape` may only appear in the last position. Actual: {shape}" + ) + + return shape + + +def convert_size(size: Size) -> Optional[StrongSize]: + """ Process a user-provided size variable into None or a valid size object. """ + if size is None: + return None + + if isinstance(size, int) or (isinstance(size, TensorVariable) and size.ndim == 0): + size = (size,) + elif isinstance(size, (list, tuple)): + size = tuple(size) + else: + raise ValueError( + f"The `size` parameter must be a tuple, TensorVariable, int or list. Actual: {type(size)}" + ) + + if isinstance(size, tuple) and Ellipsis in size: + raise ValueError(f"The `size` parameter cannot contain an Ellipsis. Actual: {size}") + + return size + + +def resize_from_dims( + dims: WeakDims, ndim_implied: int, model +) -> Tuple[int, StrongSize, StrongDims]: + """Determines a potential resize shape from a `dims` tuple. + + Parameters + ---------- + dims : array-like + A vector of dimension names, None or Ellipsis. + ndim_implied : int + Number of RV dimensions that were implied from its inputs alone. + model : pm.Model + The current model on stack. + + Returns + ------- + ndim_resize : int + Number of dimensions that should be added through resizing. + resize_shape : array-like + The shape of the new dimensions. + """ + if Ellipsis in dims: + # Auto-complete the dims tuple to the full length. + # We don't have a way to know the names of implied + # dimensions, so they will be `None`. + dims = (*dims[:-1], *[None] * ndim_implied) + + ndim_resize = len(dims) - ndim_implied + + # All resize dims must be known already (numerically or symbolically). + unknowndim_resize_dims = set(dims[:ndim_resize]) - set(model.dim_lengths) + if unknowndim_resize_dims: + raise KeyError( + f"Dimensions {unknowndim_resize_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_shape = tuple(model.dim_lengths[dname] for dname in dims[:ndim_resize]) + return ndim_resize, resize_shape, dims + + +def resize_from_observed( + observed, ndim_implied: int +) -> Tuple[int, StrongSize, Union[np.ndarray, Variable]]: + """Determines a potential resize shape from observations. + + Parameters + ---------- + observed : scalar, array-like + The value of the `observed` kwarg to the RV creation. + ndim_implied : int + Number of RV dimensions that were implied from its inputs alone. + + Returns + ------- + ndim_resize : int + Number of dimensions that should be added through resizing. + resize_shape : array-like + The shape of the new dimensions. + observed : scalar, array-like + Observations as numpy array or `Variable`. + """ + if not hasattr(observed, "shape"): + observed = pandas_to_array(observed) + ndim_resize = observed.ndim - ndim_implied + resize_shape = tuple(observed.shape[d] for d in range(ndim_resize)) + return ndim_resize, resize_shape, observed + + +def find_size(shape=None, size=None, ndim_supp=None): + """Determines the size keyword argument for creating a Distribution. + + Parameters + ---------- + shape : tuple + A tuple specifying the final shape of a distribution + size : tuple + A tuple specifying the size of a distribution + ndim_supp : int + The support dimension of the distribution. + 0 if a univariate distribution, 1 if a multivariate distribution. + + Returns + ------- + create_size : int + The size argument to be passed to the distribution + ndim_expected : int + Number of dimensions expected after distribution was created + ndim_batch : int + Number of batch dimensions + ndim_supp : int + Number of support dimensions + """ + + ndim_expected = None + ndim_batch = None + create_size = None + + if shape is not None: + if Ellipsis in shape: + # Ellipsis short-hands all implied dimensions. Therefore + # we don't know how many dimensions to expect. + ndim_expected = ndim_batch = None + # Create the RV with its implied shape and resize later + create_size = None + else: + ndim_expected = len(tuple(shape)) + ndim_batch = ndim_expected - ndim_supp + create_size = tuple(shape)[:ndim_batch] + elif size is not None: + ndim_expected = ndim_supp + len(tuple(size)) + ndim_batch = ndim_expected - ndim_supp + create_size = size + + return create_size, ndim_expected, ndim_batch, ndim_supp + + +def maybe_resize( + rv_out, + rv_op, + dist_params, + ndim_expected, + ndim_batch, + ndim_supp, + shape, + size, + **kwargs, +): + """Resize a distribution if necessary. + + Parameters + ---------- + rv_out : RandomVariable + The RandomVariable to be resized if necessary + rv_op : RandomVariable.__class__ + The RandomVariable class to recreate it + dist_params : dict + Input parameters to recreate the RandomVariable + ndim_expected : int + Number of dimensions expected after distribution was created + ndim_batch : int + Number of batch dimensions + ndim_supp : int + The support dimension of the distribution. + 0 if a univariate distribution, 1 if a multivariate distribution. + shape : tuple + A tuple specifying the final shape of a distribution + size : tuple + A tuple specifying the size of a distribution + + Returns + ------- + rv_out : int + The size argument to be passed to the distribution + """ + ndim_actual = rv_out.ndim + ndims_unexpected = ndim_actual != ndim_expected + + if shape is not None and ndims_unexpected: + if Ellipsis in shape: + # Resize and we're done! + rv_out = change_rv_size(rv_var=rv_out, new_size=shape[:-1], expand=True) + else: + # This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)). + # Recreate the RV without passing `size` to created it with just the implied dimensions. + rv_out = rv_op(*dist_params, size=None, **kwargs) + + # Now resize by any remaining "extra" dimensions that were not implied from support and parameters + if rv_out.ndim < ndim_expected: + expand_shape = shape[: ndim_expected - rv_out.ndim] + rv_out = change_rv_size(rv_var=rv_out, new_size=expand_shape, expand=True) + if not rv_out.ndim == ndim_expected: + raise ShapeError( + f"Failed to create the RV with the expected dimensionality. " + f"This indicates a severe problem. Please open an issue.", + actual=ndim_actual, + expected=ndim_batch + ndim_supp, + ) + + # Warn about the edge cases where the RV Op creates more dimensions than + # it should based on `size` and `RVOp.ndim_supp`. + if size is not None and ndims_unexpected: + warnings.warn( + f"You may have expected a ({len(tuple(size))}+{ndim_supp})-dimensional RV, but the resulting RV will be {ndim_actual}-dimensional." + ' To silence this warning use `warnings.simplefilter("ignore", pm.ShapeWarning)`.', + ShapeWarning, + ) + + return rv_out diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py index 2b75ba765ef..702693bcbd6 100644 --- a/pymc3/tests/test_shape_handling.py +++ b/pymc3/tests/test_shape_handling.py @@ -1,4 +1,4 @@ -# Copyright 2020 The PyMC Developers +# Copyright 2021 The PyMC Developers # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,15 +20,13 @@ import pymc3 as pm -from pymc3.distributions.distribution import ( - _convert_dims, - _convert_shape, - _convert_size, -) from pymc3.distributions.shape_utils import ( broadcast_dist_samples_shape, broadcast_dist_samples_to, broadcast_distribution_samples, + convert_dims, + convert_shape, + convert_size, get_broadcastable_dist_samples, shapes_broadcasting, to_tuple, @@ -419,25 +417,25 @@ def test_mvnormal_shape_size_difference(self): assert tuple(rv.shape.eval()) == (5, 4, 5, 4, 3) def test_convert_dims(self): - assert _convert_dims(dims="town") == ("town",) + assert convert_dims(dims="town") == ("town",) with pytest.raises(ValueError, match="must be a tuple, str or list"): - _convert_dims(3) + convert_dims(3) with pytest.raises(ValueError, match="may only appear in the last position"): - _convert_dims(dims=(..., "town")) + convert_dims(dims=(..., "town")) def test_convert_shape(self): - assert _convert_shape(5) == (5,) + assert convert_shape(5) == (5,) with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): - _convert_shape(shape="notashape") + convert_shape(shape="notashape") with pytest.raises(ValueError, match="may only appear in the last position"): - _convert_shape(shape=(3, ..., 2)) + convert_shape(shape=(3, ..., 2)) def test_convert_size(self): - assert _convert_size(7) == (7,) + assert convert_size(7) == (7,) with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): - _convert_size(size="notasize") + convert_size(size="notasize") with pytest.raises(ValueError, match="cannot contain"): - _convert_size(size=(3, ...)) + convert_size(size=(3, ...)) def test_lazy_flavors(self): assert pm.Uniform.dist(2, [4, 5], size=[3, 2]).eval().shape == (3, 2) From 7e747f24cb8a9de206356504b55f28f302340ba7 Mon Sep 17 00:00:00 2001 From: Michael Osthege Date: Fri, 4 Jun 2021 20:11:31 +0200 Subject: [PATCH 9/9] Use initial instead of testval internally Co-authored-by: Oriol Abril-Pla --- pymc3/distributions/distribution.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index a9f090eabed..26530af9a01 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -202,9 +202,9 @@ def __new__( ) dims = convert_dims(dims) - # Create the RV without specifying testval, because the testval may have a shape + # Create the RV without specifying initval, because the initval 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) + rv_out = cls.dist(*args, rng=rng, initval=None, **kwargs) ndim_actual = rv_out.ndim resize_shape = None