Fix Wishart shape inference when there are batched parameters

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -905,7 +905,7 @@ class WishartRV(RandomVariable):
 
     def _shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
         # The shape of second parameter `V` defines the shape of the output.
-        return dist_params[1].shape
+        return dist_params[1].shape[-2:]
 
     @classmethod
     def rng_fn(cls, rng, nu, V, size):

pymc/tests/test_distributions_random.py

@@ -1784,8 +1784,25 @@ def wishart_rng_fn(self, size, nu, V, rng):
         "check_rv_size",
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
+        "check_rv_size_batched_params",
     ]
 
+    def check_rv_size_batched_params(self):
+        for size in (None, (2,), (1, 2), (4, 3, 2)):
+            x = pm.Wishart.dist(nu=4, V=np.stack([np.eye(3), np.eye(3)]), size=size)
+
+            if size is None:
+                expected_shape = (2, 3, 3)
+            else:
+                expected_shape = size + (3, 3)
+
+            assert tuple(x.shape.eval()) == expected_shape
+
+            # RNG does not currently support batched parameters, whet it does this test
+            # should be updated to check that draws also have the expected shape
+            with pytest.raises(ValueError):
+                x.eval()
+
 
 class TestMatrixNormal(BaseTestDistributionRandom):