Fixing scalar shape handling: treat (1,) as vectors

pymc3/distributions/discrete.py

@@ -209,7 +209,7 @@ def __init__(self, alpha, beta, n, *args, **kwargs):
         self.mode = tt.cast(tround(alpha / (alpha + beta)), "int8")
 
     def _random(self, alpha, beta, n, size=None):
-        size = size or 1
+        size = size or ()
         p = stats.beta.rvs(a=alpha, b=beta, size=size).flatten()
         # Sometimes scipy.beta returns nan. Ugh.
         while np.any(np.isnan(p)):
@@ -1093,7 +1093,7 @@ def random(self, point=None, size=None):
         return generate_samples(
             random_choice,
             p=p,
-            broadcast_shape=p.shape[:-1] or (1,),
+            broadcast_shape=p.shape[:-1],
             dist_shape=self.shape,
             size=size,
         )

pymc3/distributions/distribution.py

@@ -107,6 +107,12 @@ def __new__(cls, name, *args, **kwargs):
                 dims = (dims,)
             shape = model.shape_from_dims(dims)
 
+        # failsafe against 0-shapes
+        if shape is not None and any(np.atleast_1d(shape) <= 0):
+            raise ValueError(
+                f"Distribution initialized with invalid shape {shape}. This is not allowed."
+            )
+
         # Some distributions do not accept shape=None
         if has_shape or shape is not None:
             dist = cls.dist(*args, **kwargs, shape=shape)
@@ -1002,11 +1008,9 @@ def _draw_value(param, point=None, givens=None, size=None):
 
 
 def _is_one_d(dist_shape):
-    if hasattr(dist_shape, "dshape") and dist_shape.dshape in ((), (0,), (1,)):
-        return True
-    elif hasattr(dist_shape, "shape") and dist_shape.shape in ((), (0,), (1,)):
+    if hasattr(dist_shape, "dshape") and dist_shape.dshape in {()}:
         return True
-    elif to_tuple(dist_shape) == ():
+    elif hasattr(dist_shape, "shape") and dist_shape.shape in {()}:
         return True
     return False
 
@@ -1044,6 +1048,7 @@ def generate_samples(generator, *args, **kwargs):
     Any remaining args and kwargs are passed on to the generator function.
     """
     dist_shape = kwargs.pop("dist_shape", ())
+    # TODO: the following variable is no longer used !!
     one_d = _is_one_d(dist_shape)
     size = kwargs.pop("size", None)
     broadcast_shape = kwargs.pop("broadcast_shape", None)
@@ -1120,21 +1125,5 @@ def generate_samples(generator, *args, **kwargs):
         samples = generator(size=dist_bcast_shape, *args, **kwargs)
     else:
         samples = generator(size=size_tup + dist_bcast_shape, *args, **kwargs)
-    samples = np.asarray(samples)
 
-    # reshape samples here
-    if samples.ndim > 0 and samples.shape[0] == 1 and size_tup == (1,):
-        if (
-            len(samples.shape) > len(dist_shape)
-            and samples.shape[-len(dist_shape) :] == dist_shape[-len(dist_shape) :]
-        ):
-            samples = samples.reshape(samples.shape[1:])
-
-    if (
-        one_d
-        and samples.ndim > 0
-        and samples.shape[-1] == 1
-        and (samples.shape != size_tup or size_tup == tuple() or size_tup == (1,))
-    ):
-        samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)

pymc3/distributions/posterior_predictive.py

@@ -560,8 +560,6 @@ def random_sample(
             shape: Tuple[int, ...],
         ) -> np.ndarray:
             val = meth(point=point, size=size)
-            if size == 1:
-                val = np.expand_dims(val, axis=0)
             try:
                 assert val.shape == (size,) + shape, (
                     "Sampling from random of %s yields wrong shape" % param

pymc3/tests/test_distributions_random.py

@@ -204,12 +204,19 @@ def test_blocking_context(self):
 class BaseTestCases:
     class BaseTestCase(SeededTest):
         shape = 5
+        # the following are the default values of the distribution that take effect
+        # when the parametrized shape/size in the test case is None.
+        # For every distribution that defaults to non-scalar shapes they must be
+        # specified by the inheriting Test class. example: TestGaussianRandomWalk
+        default_shape = ()
+        default_size = ()
 
         def setup_method(self, *args, **kwargs):
             super().setup_method(*args, **kwargs)
             self.model = pm.Model()
 
         def get_random_variable(self, shape, with_vector_params=False, name=None):
+            """ Creates a RandomVariable of the parametrized distribution. """
             if with_vector_params:
                 params = {
                     key: value * np.ones(self.shape, dtype=np.dtype(type(value)))
@@ -220,100 +227,87 @@ def get_random_variable(self, shape, with_vector_params=False, name=None):
             if name is None:
                 name = self.distribution.__name__
             with self.model:
-                if shape is None:
-                    return self.distribution(name, transform=None, **params)
-                else:
-                    try:
+                try:
+                    if shape is None:
+                        # in the test case parametrization "None" means "no specified (default)"
+                        return self.distribution(name, transform=None, **params)
+                    else:
                         return self.distribution(name, shape=shape, transform=None, **params)
-                    except TypeError:
-                        if np.sum(np.atleast_1d(shape)) == 0:
-                            pytest.skip("Timeseries must have positive shape")
-                        raise
+                except TypeError:
+                    if np.sum(np.atleast_1d(shape)) == 0:
+                        pytest.skip("Timeseries must have positive shape")
+                    raise
 
         @staticmethod
         def sample_random_variable(random_variable, size):
+            """ Draws samples from a RandomVariable using its .random() method. """
             try:
-                return random_variable.random(size=size)
+                if size is None:
+                    return random_variable.random()
+                else:
+                    return random_variable.random(size=size)
             except AttributeError:
-                return random_variable.distribution.random(size=size)
-
-        @pytest.mark.parametrize("size", [None, 5, (4, 5)], ids=str)
-        def test_scalar_parameter_shape(self, size):
-            rv = self.get_random_variable(None)
-            if size is None:
-                expected = (1,)
-            else:
-                expected = np.atleast_1d(size).tolist()
-            actual = np.atleast_1d(self.sample_random_variable(rv, size)).shape
-            assert tuple(expected) == actual
+                if size is None:
+                    return random_variable.distribution.random()
+                else:
+                    return random_variable.distribution.random(size=size)
 
-        @pytest.mark.parametrize("size", [None, 5, (4, 5)], ids=str)
-        def test_scalar_shape(self, size):
-            shape = 10
+        @pytest.mark.parametrize("size", [None, (), 1, (1,), 5, (4, 5)], ids=str)
+        @pytest.mark.parametrize("shape", [None, ()], ids=str)
+        def test_scalar_distribution_shape(self, shape, size):
+            """ Draws samples of different [size] from a scalar [shape] RV. """
             rv = self.get_random_variable(shape)
+            exp_shape = self.default_shape if shape is None else tuple(np.atleast_1d(shape))
+            exp_size = self.default_size if size is None else tuple(np.atleast_1d(size))
+            expected = exp_size + exp_shape
+            actual = np.shape(self.sample_random_variable(rv, size))
+            assert (
+                expected == actual
+            ), f"Sample size {size} from {shape}-shaped RV had shape {actual}. Expected: {expected}"
+            # check that negative size raises an error
+            with pytest.raises(ValueError):
+                self.sample_random_variable(rv, size=-2)
+            with pytest.raises(ValueError):
+                self.sample_random_variable(rv, size=(3, -2))
 
-            if size is None:
-                expected = []
-            else:
-                expected = np.atleast_1d(size).tolist()
-            expected.append(shape)
-            actual = np.atleast_1d(self.sample_random_variable(rv, size)).shape
-            assert tuple(expected) == actual
-
-        @pytest.mark.parametrize("size", [None, 5, (4, 5)], ids=str)
-        def test_parameters_1d_shape(self, size):
-            rv = self.get_random_variable(self.shape, with_vector_params=True)
-            if size is None:
-                expected = []
-            else:
-                expected = np.atleast_1d(size).tolist()
-            expected.append(self.shape)
-            actual = self.sample_random_variable(rv, size).shape
-            assert tuple(expected) == actual
-
-        @pytest.mark.parametrize("size", [None, 5, (4, 5)], ids=str)
-        def test_broadcast_shape(self, size):
-            broadcast_shape = (2 * self.shape, self.shape)
-            rv = self.get_random_variable(broadcast_shape, with_vector_params=True)
-            if size is None:
-                expected = []
-            else:
-                expected = np.atleast_1d(size).tolist()
-            expected.extend(broadcast_shape)
-            actual = np.atleast_1d(self.sample_random_variable(rv, size)).shape
-            assert tuple(expected) == actual
-
+        @pytest.mark.parametrize("size", [None, ()], ids=str)
         @pytest.mark.parametrize(
-            "shape", [(), (1,), (1, 1), (1, 2), (10, 10, 1), (10, 10, 2)], ids=str
+            "shape", [None, (), (1,), (1, 1), (1, 2), (10, 11, 1), (9, 10, 2)], ids=str
         )
-        def test_different_shapes_and_sample_sizes(self, shape):
-            prefix = self.distribution.__name__
-
-            rv = self.get_random_variable(shape, name=f"{prefix}_{shape}")
-            for size in (None, 1, 5, (4, 5)):
-                if size is None:
-                    s = []
-                else:
-                    try:
-                        s = list(size)
-                    except TypeError:
-                        s = [size]
-                    if s == [1]:
-                        s = []
-                if shape not in ((), (1,)):
-                    s.extend(shape)
-                e = tuple(s)
-                a = self.sample_random_variable(rv, size).shape
-                assert e == a
+        def test_scalar_sample_shape(self, shape, size):
+            """ Draws samples of scalar [size] from a [shape] RV. """
+            rv = self.get_random_variable(shape)
+            exp_shape = self.default_shape if shape is None else tuple(np.atleast_1d(shape))
+            exp_size = self.default_size if size is None else tuple(np.atleast_1d(size))
+            expected = exp_size + exp_shape
+            actual = np.shape(self.sample_random_variable(rv, size))
+            assert (
+                expected == actual
+            ), f"Sample size {size} from {shape}-shaped RV had shape {actual}. Expected: {expected}"
+
+        @pytest.mark.parametrize("size", [None, 3, (4, 5)], ids=str)
+        @pytest.mark.parametrize("shape", [None, 1, (10, 11, 1)], ids=str)
+        def test_vector_params(self, shape, size):
+            shape = self.shape
+            rv = self.get_random_variable(shape, with_vector_params=True)
+            exp_shape = self.default_shape if shape is None else tuple(np.atleast_1d(shape))
+            exp_size = self.default_size if size is None else tuple(np.atleast_1d(size))
+            expected = exp_size + exp_shape
+            actual = np.shape(self.sample_random_variable(rv, size))
+            assert (
+                expected == actual
+            ), f"Sample size {size} from {shape}-shaped RV had shape {actual}. Expected: {expected}"
+
+        @pytest.mark.parametrize("shape", [-2, 0, (0,), (2, 0), (5, 0, 3)])
+        def test_shape_error_on_zero_shape_rv(self, shape):
+            with pytest.raises(ValueError, match="not allowed"):
+                self.get_random_variable(shape)
 
 
 class TestGaussianRandomWalk(BaseTestCases.BaseTestCase):
     distribution = pm.GaussianRandomWalk
     params = {"mu": 1.0, "sigma": 1.0}
-
-    @pytest.mark.xfail(reason="Supporting this makes a nasty API")
-    def test_broadcast_shape(self):
-        super().test_broadcast_shape()
+    default_shape = (1,)
 
 
 class TestNormal(BaseTestCases.BaseTestCase):

pymc3/tests/test_shape_handling.py

@@ -214,9 +214,6 @@ def test_broadcast_dist_samples_to(self, samples_to_broadcast_to):
 def test_sample_generate_values(fixture_model, fixture_sizes):
     model, RVs = fixture_model
     size = to_tuple(fixture_sizes)
-    if size == (1,):
-        # Single draws are interpreted as scalars for backwards compatibility
-        size = tuple()
     with model:
         prior = pm.sample_prior_predictive(samples=fixture_sizes)
         for rv in RVs: