Test expected (inferred) and actual shape of draws in TestBaseDistributionRandom

* Fixes bug in returned samples from `Wishart` when `size=1`

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -891,9 +891,13 @@ def _shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
         return dist_params[1].shape
 
     @classmethod
-    def rng_fn(cls, rng, nu, V, size=None):
-        size = size if size else 1  # Default size for Scipy's wishart.rvs is 1
-        return stats.wishart.rvs(np.int(nu), V, size=size, random_state=rng)
+    def rng_fn(cls, rng, nu, V, size):
+        scipy_size = size if size else 1  # Default size for Scipy's wishart.rvs is 1
+        result = stats.wishart.rvs(np.int(nu), V, size=scipy_size, random_state=rng)
+        if size == (1,):
+            return result[np.newaxis, ...]
+        else:
+            return result
 
 
 wishart = WishartRV()

pymc/tests/test_distributions_random.py

@@ -369,8 +369,9 @@ def check_rv_size(self):
         sizes_expected = self.sizes_expected or [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)]
         for size, expected in zip(sizes_to_check, sizes_expected):
             pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size)
-            actual = tuple(pymc_rv.shape.eval())
-            assert actual == expected, f"size={size}, expected={expected}, actual={actual}"
+            expected_symbolic = tuple(pymc_rv.shape.eval())
+            actual = pymc_rv.eval().shape
+            assert actual == expected_symbolic == expected
 
         # test multi-parameters sampling for univariate distributions (with univariate inputs)
         if (
@@ -390,8 +391,9 @@ def check_rv_size(self):
             ]
             for size, expected in zip(sizes_to_check, sizes_expected):
                 pymc_rv = self.pymc_dist.dist(**params, size=size)
-                actual = tuple(pymc_rv.shape.eval())
-                assert actual == expected
+                expected_symbolic = tuple(pymc_rv.shape.eval())
+                actual = pymc_rv.eval().shape
+                assert actual == expected_symbolic == expected
 
     def validate_tests_list(self):
         assert len(self.checks_to_run) == len(
@@ -417,10 +419,18 @@ class TestFlat(BaseTestDistributionRandom):
     expected_rv_op_params = {}
     checks_to_run = [
         "check_pymc_params_match_rv_op",
-        "check_rv_size",
+        "check_rv_inferred_size",
         "check_not_implemented",
     ]
 
+    def check_rv_inferred_size(self):
+        sizes_to_check = self.sizes_to_check or [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)]
+        sizes_expected = self.sizes_expected or [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)]
+        for size, expected in zip(sizes_to_check, sizes_expected):
+            pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size)
+            expected_symbolic = tuple(pymc_rv.shape.eval())
+            assert expected_symbolic == expected
+
     def check_not_implemented(self):
         with pytest.raises(NotImplementedError):
             self.pymc_rv.eval()
@@ -432,10 +442,18 @@ class TestHalfFlat(BaseTestDistributionRandom):
     expected_rv_op_params = {}
     checks_to_run = [
         "check_pymc_params_match_rv_op",
-        "check_rv_size",
+        "check_rv_inferred_size",
         "check_not_implemented",
     ]
 
+    def check_rv_inferred_size(self):
+        sizes_to_check = self.sizes_to_check or [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)]
+        sizes_expected = self.sizes_expected or [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)]
+        for size, expected in zip(sizes_to_check, sizes_expected):
+            pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size)
+            expected_symbolic = tuple(pymc_rv.shape.eval())
+            assert expected_symbolic == expected
+
     def check_not_implemented(self):
         with pytest.raises(NotImplementedError):
             self.pymc_rv.eval()