raise ValueError when Distribution is initialized with a 0-length dimension

Author: michaelosthege

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)

pymc3/tests/test_distributions_random.py

@@ -264,6 +264,9 @@ def test_scalar_distribution_shape(self, shape, 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, match="not allowed"):
+                rv.random(size=-2)
 
         @pytest.mark.parametrize("size", [None, ()], ids=str)
         @pytest.mark.parametrize(
@@ -293,6 +296,11 @@ def test_vector_params(self, shape, size):
                 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