Implement backwards-compatible shape and dims with Ellipsis support

RELEASE-NOTES.md

@@ -10,6 +10,11 @@
 
 ### 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.
+  - 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)).
 - ...
 

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)

pymc3/distributions/distribution.py

@@ -19,22 +19,29 @@
 import warnings
 
 from abc import ABCMeta
-from typing import TYPE_CHECKING
+from typing import Optional
 
+import aesara
+import aesara.tensor as at
 import dill
 
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
 
+from pymc3.aesaraf import change_rv_size
 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.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
 
@@ -113,7 +120,52 @@ class Distribution(metaclass=DistributionMeta):
     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 +178,87 @@ 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 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, initval=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.
+
+            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
+            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 +267,38 @@ 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)
 
-        rv_var = cls.rv_op(*dist_params, rng=rng, **kwargs)
+        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)
+        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 (
-            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 +308,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"