Adapt to new size kwarg interpretation

Closes #5446

Author: ricardoV94

pymc/aesaraf.py

@@ -173,9 +173,9 @@ def change_rv_size(
     tag = rv_var.tag
 
     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(new_size) + tuple(size)
+        old_shape = tuple(rv_node.op._infer_shape(size, dist_params))
+        old_size = old_shape[: len(old_shape) - rv_node.op.ndim_supp]
+        new_size = tuple(new_size) + tuple(old_size)
 
     # 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).

pymc/distributions/multivariate.py

@@ -247,7 +247,7 @@ def dist(cls, mu, cov=None, tau=None, chol=None, lower=True, **kwargs):
     def get_moment(rv, size, mu, cov):
         moment = mu
         if not rv_size_is_none(size):
-            moment_size = at.concatenate([size, mu.shape])
+            moment_size = at.concatenate([size, [mu.shape[-1]]])
             moment = at.full(moment_size, mu)
         return moment
 
@@ -301,18 +301,13 @@ def _shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
     @classmethod
     def rng_fn(cls, rng, nu, mu, cov, size):
 
-        # Don't reassign broadcasted cov, since MvNormal expects two dimensional cov only.
-        mu, _ = broadcast_params([mu, cov], cls.ndims_params[1:])
-
-        chi2_samples = np.sqrt(rng.chisquare(nu, size=size) / nu)
-        # Add distribution shape to chi2 samples
-        chi2_samples = chi2_samples.reshape(chi2_samples.shape + (1,) * len(mu.shape))
-
         mv_samples = multivariate_normal.rng_fn(rng=rng, mean=np.zeros_like(mu), cov=cov, size=size)
 
-        size = tuple(size or ())
+        # Take chi2 draws and add an axis of length 1 to the right for correct broadcasting below
+        chi2_samples = np.sqrt(rng.chisquare(nu, size=size) / nu)[..., None]
+
         if size:
-            mu = np.broadcast_to(mu, size + mu.shape)
+            mu = np.broadcast_to(mu, size + (mu.shape[-1],))
 
         return (mv_samples / chi2_samples) + mu
 
@@ -379,7 +374,7 @@ def dist(cls, nu, Sigma=None, mu=None, cov=None, tau=None, chol=None, lower=True
     def get_moment(rv, size, nu, mu, cov):
         moment = mu
         if not rv_size_is_none(size):
-            moment_size = at.concatenate([size, moment.shape])
+            moment_size = at.concatenate([size, [mu.shape[-1]]])
             moment = at.full(moment_size, moment)
         return moment
 
@@ -453,7 +448,7 @@ def get_moment(rv, size, a):
         norm_constant = at.sum(a, axis=-1)[..., None]
         moment = a / norm_constant
         if not rv_size_is_none(size):
-            moment = at.full(at.concatenate([size, a.shape]), moment)
+            moment = at.full(at.concatenate([size, [a.shape[-1]]]), moment)
         return moment
 
     def logp(value, a):
@@ -497,8 +492,8 @@ def rng_fn(cls, rng, n, p, size):
             size = tuple(size or ())
 
             if size:
-                n = np.broadcast_to(n, size + n.shape)
-                p = np.broadcast_to(p, size + p.shape)
+                n = np.broadcast_to(n, size)
+                p = np.broadcast_to(p, size + (p.shape[-1],))
 
             res = np.empty(p.shape)
             for idx in np.ndindex(p.shape[:-1]):
@@ -571,7 +566,7 @@ def get_moment(rv, size, n, p):
         inc_bool_arr = at.abs_(diff) > 0
         mode = at.inc_subtensor(mode[inc_bool_arr.nonzero()], diff[inc_bool_arr.nonzero()])
         if not rv_size_is_none(size):
-            output_size = at.concatenate([size, p.shape])
+            output_size = at.concatenate([size, [p.shape[-1]]])
             mode = at.full(output_size, mode)
         return mode
 
@@ -623,8 +618,8 @@ def rng_fn(cls, rng, n, a, size):
             size = tuple(size or ())
 
             if size:
-                n = np.broadcast_to(n, size + n.shape)
-                a = np.broadcast_to(a, size + a.shape)
+                n = np.broadcast_to(n, size)
+                a = np.broadcast_to(a, size + (a.shape[-1],))
 
             res = np.empty(a.shape)
             for idx in np.ndindex(a.shape[:-1]):
@@ -688,10 +683,14 @@ def get_moment(rv, size, n, a):
         diff = n - at.sum(mode, axis=-1, keepdims=True)
         inc_bool_arr = at.abs_(diff) > 0
         mode = at.inc_subtensor(mode[inc_bool_arr.nonzero()], diff[inc_bool_arr.nonzero()])
-        # Reshape mode according to base shape (ignoring size)
-        mode = at.reshape(mode, rv.shape[size.size :])
+
+        # Reshape mode according to dimensions implied by the parameters
+        # This can include axes of length 1
+        _, p_bcast = broadcast_params([n, p], ndims_params=[0, 1])
+        mode = at.reshape(mode, p_bcast.shape)
+
         if not rv_size_is_none(size):
-            output_size = at.concatenate([size, mode.shape])
+            output_size = at.concatenate([size, [p.shape[-1]]])
             mode = at.full(output_size, mode)
         return mode
 
@@ -2070,7 +2069,7 @@ def make_node(self, rng, size, dtype, mu, W, alpha, tau):
         return super().make_node(rng, size, dtype, mu, W, alpha, tau)
 
     def _infer_shape(self, size, dist_params, param_shapes=None):
-        shape = tuple(size) + tuple(dist_params[0].shape)
+        shape = tuple(size) + (dist_params[0].shape[-1],)
         return shape
 
     @classmethod
@@ -2105,7 +2104,7 @@ def rng_fn(cls, rng: np.random.RandomState, mu, W, alpha, tau, size):
 
         size = tuple(size or ())
         if size:
-            mu = np.broadcast_to(mu, size + mu.shape)
+            mu = np.broadcast_to(mu, size + (mu.shape[-1],))
         z = rng.normal(size=mu.shape)
         samples = np.empty(z.shape)
         for idx in np.ndindex(mu.shape[:-1]):
@@ -2165,11 +2164,7 @@ def dist(cls, mu, W, alpha, tau, *args, **kwargs):
         return super().dist([mu, W, alpha, tau], **kwargs)
 
     def get_moment(rv, size, mu, W, alpha, tau):
-        moment = mu
-        if not rv_size_is_none(size):
-            moment_size = at.concatenate([size, moment.shape])
-            moment = at.full(moment_size, mu)
-        return moment
+        return at.full_like(rv, mu)
 
     def logp(value, mu, W, alpha, tau):
         """

pymc/distributions/shape_utils.py

@@ -26,6 +26,7 @@
 import numpy as np
 
 from aesara.graph.basic import Constant, Variable
+from aesara.graph.op import Op
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
@@ -614,10 +615,10 @@ def find_size(shape=None, size=None, ndim_supp=None):
 
 
 def maybe_resize(
-    rv_out,
-    rv_op,
+    rv_out: TensorVariable,
+    rv_op: Op,
     dist_params,
-    ndim_expected,
+    ndim_expected: int,
     ndim_batch,
     ndim_supp,
     shape,

pymc/tests/test_distributions.py

@@ -21,6 +21,7 @@
 import numpy.random as nr
 
 from aeppl.logprob import ParameterValueError
+from aesara.tensor.random.utils import broadcast_params
 
 from pymc.distributions.continuous import get_tau_sigma
 from pymc.util import UNSET
@@ -2130,9 +2131,10 @@ def test_dirichlet_invalid(self):
             (np.abs(np.random.randn(2, 2, 4)) + 1),
         ],
     )
-    @pytest.mark.parametrize("size", [2, (1, 2), (2, 4, 3)])
-    def test_dirichlet_vectorized(self, a, size):
+    @pytest.mark.parametrize("extra_size", [(2,), (1, 2), (2, 4, 3)])
+    def test_dirichlet_vectorized(self, a, extra_size):
         a = floatX(np.array(a))
+        size = extra_size + a.shape[:-1]
 
         dir = pm.Dirichlet.dist(a=a, size=size)
         vals = dir.eval()
@@ -2200,12 +2202,15 @@ def test_multinomial_p_not_normalized_symbolic(self):
             (np.abs(np.random.randn(2, 2, 4))),
         ],
     )
-    @pytest.mark.parametrize("size", [1, 2, (2, 3)])
-    def test_multinomial_vectorized(self, n, p, size):
+    @pytest.mark.parametrize("extra_size", [(1,), (2,), (2, 3)])
+    def test_multinomial_vectorized(self, n, p, extra_size):
         n = intX(np.array(n))
         p = floatX(np.array(p))
         p /= p.sum(axis=-1, keepdims=True)
 
+        _, bcast_p = broadcast_params([n, p], ndims_params=[0, 1])
+        size = extra_size + bcast_p.shape[:-1]
+
         mn = pm.Multinomial.dist(n=n, p=p, size=size)
         vals = mn.eval()
 
@@ -2269,11 +2274,14 @@ def test_dirichlet_multinomial_matches_beta_binomial(self):
             (np.abs(np.random.randn(2, 2, 4))),
         ],
     )
-    @pytest.mark.parametrize("size", [1, 2, (2, 3)])
-    def test_dirichlet_multinomial_vectorized(self, n, a, size):
+    @pytest.mark.parametrize("extra_size", [(1,), (2,), (2, 3)])
+    def test_dirichlet_multinomial_vectorized(self, n, a, extra_size):
         n = intX(np.array(n))
         a = floatX(np.array(a))
 
+        _, bcast_a = broadcast_params([n, a], ndims_params=[0, 1])
+        size = extra_size + bcast_a.shape[:-1]
+
         dm = pm.DirichletMultinomial.dist(n=n, a=a, size=size)
         vals = dm.eval()
 

pymc/tests/test_distributions_moments.py

@@ -789,7 +789,7 @@ def test_discrete_weibull_moment(q, beta, size, expected):
         ),
         (
             np.array([[1, 2, 3], [5, 6, 7]]),
-            7,
+            (7, 2),
             np.apply_along_axis(
                 lambda x: np.divide(x, np.array([6, 18])),
                 1,
@@ -798,7 +798,7 @@ def test_discrete_weibull_moment(q, beta, size, expected):
         ),
         (
             np.full(shape=np.array([7, 3]), fill_value=np.array([13, 17, 19])),
-            (11, 5),
+            (11, 5, 7),
             np.broadcast_to([13, 17, 19], shape=[11, 5, 7, 3]) / 49,
         ),
     ],
@@ -940,7 +940,7 @@ def test_interpolated_moment(x_points, pdf_points, size, expected):
         (
             np.array([[3.0, 5], [1, 4]]),
             np.identity(2),
-            (4, 5),
+            (4, 5, 2),
             np.full((4, 5, 2, 2), [[3.0, 5], [1, 4]]),
         ),
     ],
@@ -965,7 +965,7 @@ def test_mv_normal_moment(mu, cov, size, expected):
         (np.array([1, 0, 3.0, 4]), (5, 3), np.full((5, 3, 4), [1, 0, 3.0, 4])),
         (
             np.array([[3.0, 5, 2, 1], [1, 4, 0.5, 9]]),
-            (4, 5),
+            (4, 5, 2),
             np.full((4, 5, 2, 4), [[3.0, 5, 2, 1], [1, 4, 0.5, 9]]),
         ),
     ],
@@ -1008,8 +1008,8 @@ def test_moyal_moment(mu, sigma, size, expected):
         (2, rand1d, np.eye(2), 2, np.full((2, 2), rand1d)),
         (2, rand1d, np.eye(2), (2, 5), np.full((2, 5, 2), rand1d)),
         (2, rand2d, np.eye(3), None, rand2d),
-        (2, rand2d, np.eye(3), 2, np.full((2, 2, 3), rand2d)),
-        (2, rand2d, np.eye(3), (2, 5), np.full((2, 5, 2, 3), rand2d)),
+        (2, rand2d, np.eye(3), (2, 2), np.full((2, 2, 3), rand2d)),
+        (2, rand2d, np.eye(3), (2, 5, 2), np.full((2, 5, 2, 3), rand2d)),
     ],
 )
 def test_mvstudentt_moment(nu, mu, cov, size, expected):
@@ -1326,7 +1326,7 @@ def test_polyagamma_moment(h, z, size, expected):
         (
             np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]),
             np.array([1, 10]),
-            2,
+            (2, 2),
             np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]]),
         ),
     ],
@@ -1476,7 +1476,7 @@ def test_lkjcholeskycov_moment(n, eta, size, expected):
         (
             np.array([[26, 26, 26, 22]]),  # Dim: 1 x 4
             np.array([[1], [10]]),  # Dim: 2 x 1
-            (2, 1),
+            (2, 1, 2, 1),
             np.full(
                 (2, 1, 2, 1, 4),
                 np.array([[[1, 0, 0, 0]], [[2, 3, 3, 2]]]),  # Dim: 2 x 1 x 4

pymc/tests/test_distributions_random.py

@@ -1197,8 +1197,8 @@ def test_issue_3706(self):
 
 class TestMvStudentTCov(BaseTestDistributionRandom):
     def mvstudentt_rng_fn(self, size, nu, mu, cov, rng):
-        chi2_samples = rng.chisquare(nu, size=size)
         mv_samples = rng.multivariate_normal(np.zeros_like(mu), cov, size=size)
+        chi2_samples = rng.chisquare(nu, size=size)
         return (mv_samples / np.sqrt(chi2_samples[:, None] / nu)) + mu
 
     pymc_dist = pm.MvStudentT
@@ -1345,7 +1345,7 @@ def check_random_draws(self):
         draws = pm.DirichletMultinomial.dist(
             n=np.array([5, 100]),
             a=np.array([[0.001, 0.001, 0.001, 1000], [1000, 1000, 0.001, 0.001]]),
-            size=(2, 3),
+            size=(2, 3, 2),
             rng=default_rng,
         ).eval()
         assert np.all(draws.sum(-1) == np.array([5, 100]))
@@ -1361,7 +1361,7 @@ class TestDirichletMultinomial_1D_n_2D_a(BaseTestDistributionRandom):
         "n": np.array([23, 29]),
         "a": np.array([[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]),
     }
-    sizes_to_check = [None, 1, (4,), (3, 4)]
+    sizes_to_check = [None, (1, 2), (4, 2), (3, 4, 2)]
     sizes_expected = [(2, 4), (1, 2, 4), (4, 2, 4), (3, 4, 2, 4)]
     checks_to_run = ["check_rv_size"]
 

pymc/tests/test_shape_handling.py

@@ -31,7 +31,6 @@
     shapes_broadcasting,
     to_tuple,
 )
-from pymc.exceptions import ShapeWarning
 
 test_shapes = [
     (tuple(), (1,), (4,), (5, 4)),
@@ -323,7 +322,7 @@ def test_simultaneous_size_and_dims(self, with_dims_ellipsis):
             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, mu=at.ones((3, 4)), cov=at.eye(4))
+            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)
@@ -434,17 +433,11 @@ def test_mvnormal_shape_size_difference(self):
         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)
+        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)
+        rv = pm.MvNormal.dist(mu=np.ones((5, 4, 3)), cov=np.eye(3), size=(5, 4))
+        assert tuple(rv.shape.eval()) == (5, 4, 3)
 
     def test_convert_dims(self):
         assert convert_dims(dims="town") == ("town",)