Remove ellipsis functionality in shape and dims

Author: ricardoV94

pymc/distributions/distribution.py

@@ -19,7 +19,7 @@
 
 from abc import ABCMeta
 from functools import singledispatch
-from typing import Callable, Optional, Sequence, Tuple, Union, cast
+from typing import Callable, Optional, Sequence, Tuple, Union
 
 import aesara
 import numpy as np
@@ -39,7 +39,6 @@
     Shape,
     Size,
     StrongShape,
-    WeakDims,
     convert_dims,
     convert_shape,
     convert_size,
@@ -158,7 +157,7 @@ def _make_rv_and_resize_shape(
     observed,
     args,
     **kwargs,
-) -> Tuple[Variable, Optional[WeakDims], Optional[Union[np.ndarray, Variable]], StrongShape]:
+) -> Tuple[Variable, Optional[Dims], Optional[Union[np.ndarray, Variable]], StrongShape]:
     """Creates the RV and processes dims or observed to determine a resize shape."""
     # Create the RV without dims information, because that's not something tracked at the Aesara level.
     # If necessary we'll later replicate to a different size implied by already known dims.
@@ -173,9 +172,6 @@ def _make_rv_and_resize_shape(
     if dims is not None:
         if dims_can_resize:
             resize_shape, dims = resize_from_dims(dims, ndim_actual, model)
-        elif Ellipsis in dims:
-            # Replace ... with None entries to match the actual dimensionality.
-            dims = (*dims[:-1], *[None] * ndim_actual)[:ndim_actual]
     elif observed is not None:
         resize_shape, observed = resize_from_observed(observed, ndim_actual)
     return rv_out, dims, observed, resize_shape
@@ -305,9 +301,6 @@ 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.
         **kwargs
             Keyword arguments that will be forwarded to the Aesara RV Op.
             Most prominently: ``size`` or ``dtype``.
@@ -350,11 +343,6 @@ def dist(
         # This is not necessarily the final result.
         rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
 
-        # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
-        if shape is not None and Ellipsis in shape:
-            replicate_shape = cast(StrongShape, shape[:-1])
-            rv_out = change_rv_size(rv=rv_out, new_size=replicate_shape, expand=True)
-
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
         rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
         rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
@@ -508,8 +496,6 @@ 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.
 
@@ -550,11 +536,6 @@ def dist(
         # This is not necessarily the final result.
         graph = cls.rv_op(*dist_params, size=create_size, **kwargs)
 
-        # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
-        if shape is not None and Ellipsis in shape:
-            replicate_shape = cast(StrongShape, shape[:-1])
-            graph = cls.change_size(rv=graph, new_size=replicate_shape, expand=True)
-
         # TODO: Create new attr error stating that these are not available for DerivedDistribution
         # rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
         # rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")

pymc/distributions/shape_utils.py

@@ -18,7 +18,7 @@
 samples from probability distributions for stochastic nodes in PyMC.
 """
 
-from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union, cast
+from typing import Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
@@ -409,34 +409,18 @@ def broadcast_dist_samples_to(to_shape, samples, size=None):
     return [np.broadcast_to(o, to_shape) for o in samples]
 
 
-# Workaround to annotate the Ellipsis type, posted by the BDFL himself.
-# See https://github.com/python/typing/issues/684#issuecomment-548203158
-if TYPE_CHECKING:
-    from enum import Enum
-
-    class ellipsis(Enum):
-        Ellipsis = "..."
-
-    Ellipsis = ellipsis.Ellipsis
-else:
-    ellipsis = type(Ellipsis)
-
 # User-provided can be lazily specified as scalars
-Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable, ellipsis]]]
-Dims: TypeAlias = Union[str, Sequence[Optional[Union[str, ellipsis]]]]
+Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable]]]
+Dims: TypeAlias = Union[str, Sequence[Optional[str]]]
 Size: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable]]]
 
 # After conversion to vectors
-WeakShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable, ellipsis], ...]]
-WeakDims: TypeAlias = Tuple[Optional[Union[str, ellipsis]], ...]
-
-# After Ellipsis were substituted
 StrongShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]]
 StrongDims: TypeAlias = Sequence[Optional[str]]
 StrongSize: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]]
 
 
-def convert_dims(dims: Optional[Dims]) -> Optional[WeakDims]:
+def convert_dims(dims: Optional[Dims]) -> Optional[StrongDims]:
     """Process a user-provided dims variable into None or a valid dims tuple."""
     if dims is None:
         return None
@@ -448,13 +432,10 @@ def convert_dims(dims: Optional[Dims]) -> Optional[WeakDims]:
     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]:
+def convert_shape(shape: Shape) -> Optional[StrongShape]:
     """Process a user-provided shape variable into None or a valid shape object."""
     if shape is None:
         return None
@@ -468,10 +449,6 @@ def convert_shape(shape: Shape) -> Optional[WeakShape]:
         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
 
@@ -490,19 +467,17 @@ def convert_size(size: Size) -> Optional[StrongSize]:
         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[StrongSize, StrongDims]:
+def resize_from_dims(dims: Dims, ndim_implied: int, model) -> Tuple[StrongSize, StrongDims]:
     """Determines a potential resize shape from a `dims` tuple.
 
     Parameters
     ----------
     dims : array-like
-        A vector of dimension names, None or Ellipsis.
+        A vector of dimension names or None.
     ndim_implied : int
         Number of RV dimensions that were implied from its inputs alone.
     model : pm.Model
@@ -515,11 +490,6 @@ def resize_from_dims(dims: WeakDims, ndim_implied: int, model) -> Tuple[StrongSi
     dims : tuple of (str or None)
         Names or None for all dimensions after resizing.
     """
-    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)
     sdims = cast(StrongDims, dims)
 
     ndim_resize = len(sdims) - ndim_implied
@@ -565,7 +535,7 @@ def resize_from_observed(
 
 
 def find_size(
-    shape: Optional[WeakShape],
+    shape: Optional[StrongShape],
     size: Optional[StrongSize],
     ndim_supp: int,
 ) -> Tuple[Optional[StrongSize], Optional[int], Optional[int], int]:
@@ -598,16 +568,9 @@ def find_size(
     create_size: Optional[StrongSize] = 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]
+        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

pymc/tests/test_shape_handling.py

@@ -211,9 +211,7 @@ class TestShapeDimsSize:
         [
             "implicit",
             "shape",
-            "shape...",
             "dims",
-            "dims...",
             "size",
         ],
     )
@@ -249,65 +247,36 @@ 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 == "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.")
 
-    @pytest.mark.parametrize("ellipsis_in", ["none", "shape", "dims", "both"])
-    def test_simultaneous_shape_and_dims(self, ellipsis_in):
+    def test_simultaneous_shape_and_dims(self):
         with pm.Model() as pmodel:
             x = pm.ConstantData("x", [1, 2, 3], dims="ddata")
 
-            if ellipsis_in == "none":
-                # The shape and dims tuples correspond to each other.
-                # Note: No checks are performed that implied shape (x), shape and dims actually match.
-                y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", "ddata"))
-                assert pmodel.RV_dims["y"] == ("dshape", "ddata")
-            elif ellipsis_in == "shape":
-                y = pm.Normal("y", mu=x, shape=(2, ...), dims=("dshape", "ddata"))
-                assert pmodel.RV_dims["y"] == ("dshape", "ddata")
-            elif ellipsis_in == "dims":
-                y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", ...))
-                assert pmodel.RV_dims["y"] == ("dshape", None)
-            elif ellipsis_in == "both":
-                y = pm.Normal("y", mu=x, shape=(2, ...), dims=("dshape", ...))
-                assert pmodel.RV_dims["y"] == ("dshape", None)
+            # The shape and dims tuples correspond to each other.
+            # Note: No checks are performed that implied shape (x), shape and dims actually match.
+            y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", "ddata"))
+            assert pmodel.RV_dims["y"] == ("dshape", "ddata")
 
             assert "dshape" in pmodel.dim_lengths
             assert y.eval().shape == (2, 3)
 
-    @pytest.mark.parametrize("with_dims_ellipsis", [False, True])
-    def test_simultaneous_size_and_dims(self, with_dims_ellipsis):
+    def test_simultaneous_size_and_dims(self):
         with pm.Model() as pmodel:
             x = pm.ConstantData("x", [1, 2, 3], dims="ddata")
             assert "ddata" in pmodel.dim_lengths
 
             # Size does not include support dims, so this test must use a dist with support dims.
             kwargs = dict(name="y", size=(2, 3), mu=at.ones((3, 4)), cov=at.eye(4))
-            if with_dims_ellipsis:
-                y = pm.MvNormal(**kwargs, dims=("dsize", ...))
-                assert pmodel.RV_dims["y"] == ("dsize", None, None)
-            else:
-                y = pm.MvNormal(**kwargs, dims=("dsize", "ddata", "dsupport"))
-                assert pmodel.RV_dims["y"] == ("dsize", "ddata", "dsupport")
+            y = pm.MvNormal(**kwargs, dims=("dsize", "ddata", "dsupport"))
+            assert pmodel.RV_dims["y"] == ("dsize", "ddata", "dsupport")
 
             assert "dsize" in pmodel.dim_lengths
             assert y.eval().shape == (2, 3, 4)
@@ -382,7 +351,6 @@ 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, size=(3,)).eval().shape == (3,)
         assert pm.Normal.dist(mu=mu, size=(4, 3)).eval().shape == (4, 3)
 
@@ -408,10 +376,6 @@ 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=[1, 2, 3], cov=np.eye(3), size=(5, 4))
         assert tuple(rv.shape.eval()) == (5, 4, 3)
 
@@ -422,22 +386,16 @@ 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)