Repository: pymc-devs/pymc

Issue:
Reconsider not allowing shape/ellipsis in `dist` classmethods
When working in #5690, I realized we have a bug when doing the following:

```python
import pymc as pm
x = pm.GaussianRandomWalk.dist(init=pm.Normal.dist(), steps=5, shape=(3, ...))
print(x.shape.eval())  # [3, 3, 6]; should be [3, 6]
print(x.eval())  # ValueError
```

This happens because we have to resize `init` to be of the same size of `x`, as the GaussianRandomWalk (GRW) is a composite distribution of `init` + `innovations`. We do this in `dist`, and I simply did not consider that shape could have an ellipsis. 

In general, this sort of issue will be prevalent in most "composite" distributions, such as `LKJCorr`, `Censored`, `Mixture`, and so on, where we need to resize some of the the `dist` arguments to match `size, shape, dims` or else ask users to to it themselves. We can make things quite simpler for us if we don't allow the ellipsis or shape argument into the `dist` classmethod, as we will only have to worry about `size` in that case.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5719?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5719](https://codecov.io/gh/pymc-devs/pymc/pull/5719?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (4c7384b) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/a626feab2683a922fbcc6468f3cb61ab22415ab3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (a626fea) will **decrease** coverage by `14.85%`.
> The diff coverage is `15.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5719/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5719?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5719       +/-   ##
===========================================
- Coverage   88.92%   74.07%   -14.86%     
===========================================
  Files          75       75               
  Lines       13745    13740        -5     
===========================================
- Hits        12223    10178     -2045     
- Misses       1522     3562     +2040     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5719?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/distributions/timeseries.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL3RpbWVzZXJpZXMucHk=) | `28.41% <0.00%> (-18.90%)` | :arrow_down: |
| [pymc/distributions/censored.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2NlbnNvcmVkLnB5) | `40.35% <14.28%> (-54.65%)` | :arrow_down: |
| [pymc/distributions/multivariate.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL211bHRpdmFyaWF0ZS5weQ==) | `41.80% <28.57%> (-50.36%)` | :arrow_down: |
| [pymc/model\_graph.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9tb2RlbF9ncmFwaC5weQ==) | `15.86% <0.00%> (-69.66%)` | :arrow_down: |
| [pymc/func\_utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9mdW5jX3V0aWxzLnB5) | `22.50% <0.00%> (-65.00%)` | :arrow_down: |
| [pymc/printing.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9wcmludGluZy5weQ==) | `21.78% <0.00%> (-64.36%)` | :arrow_down: |
| [pymc/distributions/bound.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2JvdW5kLnB5) | `45.54% <0.00%> (-54.46%)` | :arrow_down: |
| [pymc/variational/updates.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC91cGRhdGVzLnB5) | `37.93% <0.00%> (-54.19%)` | :arrow_down: |
| [pymc/distributions/shape\_utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL3NoYXBlX3V0aWxzLnB5) | `48.53% <0.00%> (-48.54%)` | :arrow_down: |
| ... and [24 more](https://codecov.io/gh/pymc-devs/pymc/pull/5719/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |


Patch:
diff --git a/pymc/distributions/censored.py b/pymc/distributions/censored.py
index 60edd7e22..c0790dee3 100644
--- a/pymc/distributions/censored.py
+++ b/pymc/distributions/censored.py
@@ -18,6 +18,7 @@ from aesara.scalar import Clip
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
 
+from pymc.aesaraf import change_rv_size
 from pymc.distributions.distribution import SymbolicDistribution, _moment
 from pymc.util import check_dist_not_registered
 
@@ -74,10 +75,13 @@ class Censored(SymbolicDistribution):
 
     @classmethod
     def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
-        if lower is None:
-            lower = at.constant(-np.inf)
-        if upper is None:
-            upper = at.constant(np.inf)
+
+        lower = at.constant(-np.inf) if lower is None else at.as_tensor_variable(lower)
+        upper = at.constant(np.inf) if upper is None else at.as_tensor_variable(upper)
+
+        # When size is not specified, dist may have to be broadcasted according to lower/upper
+        dist_shape = size if size is not None else at.broadcast_shape(dist, lower, upper)
+        dist = change_rv_size(dist, dist_shape)
 
         # Censoring is achieved by clipping the base distribution between lower and upper
         rv_out = at.clip(dist, lower, upper)
@@ -88,8 +92,6 @@ class Censored(SymbolicDistribution):
         rv_out.tag.lower = lower
         rv_out.tag.upper = upper
 
-        if size is not None:
-            rv_out = cls.change_size(rv_out, size)
         if rngs is not None:
             rv_out = cls.change_rngs(rv_out, rngs)
 
@@ -101,12 +103,10 @@ class Censored(SymbolicDistribution):
 
     @classmethod
     def change_size(cls, rv, new_size, expand=False):
-        dist_node = rv.tag.dist.owner
+        dist = rv.tag.dist
         lower = rv.tag.lower
         upper = rv.tag.upper
-        rng, old_size, dtype, *dist_params = dist_node.inputs
-        new_size = new_size if not expand else tuple(new_size) + tuple(old_size)
-        new_dist = dist_node.op.make_node(rng, new_size, dtype, *dist_params).default_output()
+        new_dist = change_rv_size(dist, new_size, expand=expand)
         return cls.rv_op(new_dist, lower, upper)
 
     @classmethod
diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py
index 354061bab..940b3907e 100644
--- a/pymc/distributions/multivariate.py
+++ b/pymc/distributions/multivariate.py
@@ -32,7 +32,7 @@ from aesara.tensor import gammaln, sigmoid
 from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace
 from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal
 from aesara.tensor.random.op import RandomVariable, default_supp_shape_from_params
-from aesara.tensor.random.utils import broadcast_params
+from aesara.tensor.random.utils import broadcast_params, normalize_size_param
 from aesara.tensor.slinalg import Cholesky
 from aesara.tensor.slinalg import solve_lower_triangular as solve_lower
 from aesara.tensor.slinalg import solve_upper_triangular as solve_upper
@@ -1134,6 +1134,19 @@ class _LKJCholeskyCovRV(RandomVariable):
 
         D = at.as_tensor_variable(D)
 
+        # We resize the sd_dist `D` automatically so that it has (size x n) independent
+        # draws which is what the `_LKJCholeskyCovRV.rng_fn` expects. This makes the
+        # random and logp methods equivalent, as the latter also assumes a unique value
+        # for each diagonal element.
+        # Since `eta` and `n` are forced to be scalars we don't need to worry about
+        # implied batched dimensions for the time being.
+        size = normalize_size_param(size)
+        if D.owner.op.ndim_supp == 0:
+            D = change_rv_size(D, at.concatenate((size, (n,))))
+        else:
+            # The support shape must be `n` but we have no way of controlling it
+            D = change_rv_size(D, size)
+
         return super().make_node(rng, size, dtype, n, eta, D)
 
     def _infer_shape(self, size, dist_params, param_shapes=None):
@@ -1179,7 +1192,7 @@ class _LKJCholeskyCov(Continuous):
         return super().__new__(cls, name, eta, n, sd_dist, **kwargs)
 
     @classmethod
-    def dist(cls, eta, n, sd_dist, size=None, **kwargs):
+    def dist(cls, eta, n, sd_dist, **kwargs):
         eta = at.as_tensor_variable(floatX(eta))
         n = at.as_tensor_variable(intX(n))
 
@@ -1191,18 +1204,6 @@ class _LKJCholeskyCov(Continuous):
         ):
             raise TypeError("sd_dist must be a scalar or vector distribution variable")
 
-        # We resize the sd_dist automatically so that it has (size x n) independent draws
-        # which is what the `_LKJCholeskyCovRV.rng_fn` expects. This makes the random
-        # and logp methods equivalent, as the latter also assumes a unique value for each
-        # diagonal element.
-        # Since `eta` and `n` are forced to be scalars we don't need to worry about
-        # implied batched dimensions for the time being.
-        if sd_dist.owner.op.ndim_supp == 0:
-            sd_dist = change_rv_size(sd_dist, to_tuple(size) + (n,))
-        else:
-            # The support shape must be `n` but we have no way of controlling it
-            sd_dist = change_rv_size(sd_dist, to_tuple(size))
-
         # sd_dist is part of the generative graph, but should be completely ignored
         # by the logp graph, since the LKJ logp explicitly includes these terms.
         # Setting sd_dist.tag.ignore_logprob to True, will prevent Aeppl warning about
@@ -1211,7 +1212,7 @@ class _LKJCholeskyCov(Continuous):
         #  sd_dist prior components from the logp expression.
         sd_dist.tag.ignore_logprob = True
 
-        return super().dist([n, eta, sd_dist], size=size, **kwargs)
+        return super().dist([n, eta, sd_dist], **kwargs)
 
     def moment(rv, size, n, eta, sd_dists):
         diag_idxs = (at.cumsum(at.arange(1, n + 1)) - 1).astype("int32")
diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py
index bf6addb90..c7225ad50 100644
--- a/pymc/distributions/timeseries.py
+++ b/pymc/distributions/timeseries.py
@@ -19,12 +19,13 @@ import numpy as np
 
 from aesara import scan
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.utils import normalize_size_param
 
 from pymc.aesaraf import change_rv_size, floatX, intX
 from pymc.distributions import distribution, logprob, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
-from pymc.distributions.shape_utils import to_tuple
+from pymc.distributions.shape_utils import rv_size_is_none, to_tuple
 from pymc.util import check_dist_not_registered
 
 __all__ = [
@@ -54,6 +55,16 @@ class GaussianRandomWalkRV(RandomVariable):
         if not steps.ndim == 0 or not steps.dtype.startswith("int"):
             raise ValueError("steps must be an integer scalar (ndim=0).")
 
+        mu = at.as_tensor_variable(mu)
+        sigma = at.as_tensor_variable(sigma)
+        init = at.as_tensor_variable(init)
+
+        # Resize init distribution
+        size = normalize_size_param(size)
+        # If not explicit, size is determined by the shapes of mu, sigma, and init
+        init_size = size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init)
+        init = change_rv_size(init, init_size)
+
         return super().make_node(rng, size, dtype, mu, sigma, init, steps)
 
     def _supp_shape_from_params(self, dist_params, reop_param_idx=0, param_shapes=None):
@@ -116,7 +127,6 @@ class GaussianRandomWalkRV(RandomVariable):
 
         # If size is None then the returned series should be (*size, 1+steps)
         else:
-            init_size = (*size, 1)
             dist_shape = (*size, int(steps))
 
         innovations = rng.normal(loc=mu, scale=sigma, size=dist_shape)
@@ -161,15 +171,9 @@ class GaussianRandomWalk(distribution.Continuous):
             raise ValueError("Must specify steps parameter")
         steps = at.as_tensor_variable(intX(steps))
 
-        shape = kwargs.get("shape", None)
-        if size is None and shape is None:
-            init_size = None
-        else:
-            init_size = to_tuple(size) if size is not None else to_tuple(shape)[:-1]
-
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
         if init is None:
-            init = Normal.dist(mu, sigma, size=init_size)
+            init = Normal.dist(mu, sigma)
         else:
             if not (
                 isinstance(init, at.TensorVariable)
@@ -179,13 +183,6 @@ class GaussianRandomWalk(distribution.Continuous):
             ):
                 raise TypeError("init must be a univariate distribution variable")
 
-            if init_size is not None:
-                init = change_rv_size(init, init_size)
-            else:
-                # If not explicit, size is determined by the shapes of mu, sigma, and init
-                bcast_shape = at.broadcast_arrays(mu, sigma, init)[0].shape
-                init = change_rv_size(init, bcast_shape)
-
         # Ignores logprob of init var because that's accounted for in the logp method
         init.tag.ignore_logprob = True
 


Test patch:
diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py
index 661b05d91..cb485a086 100644
--- a/pymc/tests/test_distributions.py
+++ b/pymc/tests/test_distributions.py
@@ -3381,6 +3381,18 @@ class TestCensored:
         new_dist = pm.Censored.change_size(base_dist, (4,), expand=True)
         assert new_dist.eval().shape == (4, 3, 2)
 
+    def test_dist_broadcasted_by_lower_upper(self):
+        x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=None)
+        assert tuple(x.owner.inputs[0].shape.eval()) == (2,)
+
+        x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=np.zeros((4, 2)))
+        assert tuple(x.owner.inputs[0].shape.eval()) == (4, 2)
+
+        x = pm.Censored.dist(
+            pm.Normal.dist(size=(3, 4, 2)), lower=np.zeros((2,)), upper=np.zeros((4, 2))
+        )
+        assert tuple(x.owner.inputs[0].shape.eval()) == (3, 4, 2)
+
 
 class TestLKJCholeskCov:
     def test_dist(self):
@@ -3425,8 +3437,18 @@ class TestLKJCholeskCov:
             pm.MvNormal.dist(np.ones(3), np.eye(3)),
         ],
     )
-    def test_sd_dist_automatically_resized(self, sd_dist):
-        x = pm.LKJCholeskyCov.dist(n=3, eta=1, sd_dist=sd_dist, size=10, compute_corr=False)
+    @pytest.mark.parametrize(
+        "size, shape",
+        [
+            ((10,), None),
+            (None, (10, 6)),
+            (None, (10, ...)),
+        ],
+    )
+    def test_sd_dist_automatically_resized(self, sd_dist, size, shape):
+        x = pm.LKJCholeskyCov.dist(
+            n=3, eta=1, sd_dist=sd_dist, size=size, shape=shape, compute_corr=False
+        )
         resized_sd_dist = x.owner.inputs[-1]
         assert resized_sd_dist.eval().shape == (10, 3)
         # LKJCov has support shape `(n * (n+1)) // 2`
diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py
index 76bb91dae..b0c5325bc 100644
--- a/pymc/tests/test_distributions_timeseries.py
+++ b/pymc/tests/test_distributions_timeseries.py
@@ -103,6 +103,12 @@ def test_gaussian_random_walk_init_dist_shape(init):
     assert tuple(grw.owner.inputs[-2].shape.eval()) == (3, 2)
 
 
+def test_shape_ellipsis():
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=5, init=pm.Normal.dist(), shape=(3, ...))
+    assert tuple(grw.shape.eval()) == (3, 6)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (3,)
+
+
 def test_gaussianrandomwalk_broadcasted_by_init_dist():
     grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=4, init=pm.Normal.dist(size=(2, 3)))
     assert tuple(grw.shape.eval()) == (2, 3, 5)
diff --git a/pymc/tests/test_sampling.py b/pymc/tests/test_sampling.py
index 319b656a9..8e3550021 100644
--- a/pymc/tests/test_sampling.py
+++ b/pymc/tests/test_sampling.py
@@ -1406,8 +1406,8 @@ class TestDraw(SeededTest):
         assert np.all(draws == np.arange(5))
 
 
-class test_step_args(SeededTest):
-    with pm.Model() as model:
+def test_step_args():
+    with pm.Model(rng_seeder=1410) as model:
         a = pm.Normal("a")
         idata0 = pm.sample(target_accept=0.5)
         idata1 = pm.sample(nuts={"target_accept": 0.5})
@@ -1415,7 +1415,7 @@ class test_step_args(SeededTest):
     npt.assert_almost_equal(idata0.sample_stats.acceptance_rate.mean(), 0.5, decimal=1)
     npt.assert_almost_equal(idata1.sample_stats.acceptance_rate.mean(), 0.5, decimal=1)
 
-    with pm.Model() as model:
+    with pm.Model(rng_seeder=1418) as model:
         a = pm.Normal("a")
         b = pm.Poisson("b", 1)
         idata0 = pm.sample(target_accept=0.5)