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

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

Author: michaelosthege

RELEASE-NOTES.md

@@ -1,7 +1,19 @@
 # Release Notes
 
 ## PyMC3 vNext (on deck)
-This is the first release to support Python3.9 and to drop Python3.6.
+This release breaks some APIs w.r.t. `3.10.0`.
+It also brings some dreadfully awaited fixes, so be sure to go through the changes below.
+(Or latest when you run into problems.)
+
+### Breaking Changes
+- Python 3.6 support was dropped (by no longer testing) and Python 3.9 was added (see [#4332](https://github.com/pymc-devs/pymc3/pull/4332)).
+- Changed shape behavior: __No longer collapse length 1 vector shape into scalars.__ (see [#4206](https://github.com/pymc-devs/pymc3/issue/4206) and [#4214](https://github.com/pymc-devs/pymc3/pull/4214))
+  - __Applies to random variables and also the `.random(size=...)` kwarg!__
+  - To create scalar variables you must now use `shape=None` or `shape=()`.
+  - __`shape=(1,)` and `shape=1` now become vectors.__ Previously they were collapsed into scalars
+  - 0-length dimensions are now ruled illegal for random variables and raise a `ValueError`.
+- In `sample_prior_predictive` the `vars` kwarg was removed in favor of `var_names` (see [#4327](https://github.com/pymc-devs/pymc3/pull/4327)).
+- Removed `theanof.set_theano_config` because it illegally changed Theano's internal state (see [#4329](https://github.com/pymc-devs/pymc3/pull/4329)).
 
 ### New Features
 - `OrderedProbit` distribution added (see [#4232](https://github.com/pymc-devs/pymc3/pull/4232)).
@@ -10,8 +22,6 @@ This is the first release to support Python3.9 and to drop Python3.6.
 ### Maintenance
 - Fixed bug whereby partial traces returns after keyboard interrupt during parallel sampling had fewer draws than would've been available [#4318](https://github.com/pymc-devs/pymc3/pull/4318)
 - Make `sample_shape` same across all contexts in `draw_values` (see [#4305](https://github.com/pymc-devs/pymc3/pull/4305)).
-- Removed `theanof.set_theano_config` because it illegally touched Theano's privates (see [#4329](https://github.com/pymc-devs/pymc3/pull/4329)).
-- In `sample_posterior_predictive` the `vars` kwarg was removed in favor of `var_names` (see [#4343](https://github.com/pymc-devs/pymc3/pull/4343)).
 - The notebook gallery has been moved to https://github.com/pymc-devs/pymc-examples (see [#4348](https://github.com/pymc-devs/pymc3/pull/4348)).
 - `math.logsumexp` now matches `scipy.special.logsumexp` when arrays contain infinite values (see [#4360](https://github.com/pymc-devs/pymc3/pull/4360)).
 - Fixed mathematical formulation in `MvStudentT` random method. (see [#4359](https://github.com/pymc-devs/pymc3/pull/4359))

pymc3/distributions/discrete.py

@@ -212,7 +212,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)):
@@ -1096,7 +1096,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)
@@ -1001,16 +1007,6 @@ def _draw_value(param, point=None, givens=None, size=None):
     raise ValueError("Unexpected type in draw_value: %s" % type(param))
 
 
-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,)):
-        return True
-    elif to_tuple(dist_shape) == ():
-        return True
-    return False
-
-
 def generate_samples(generator, *args, **kwargs):
     """Generate samples from the distribution of a random variable.
 
@@ -1044,7 +1040,6 @@ def generate_samples(generator, *args, **kwargs):
     Any remaining args and kwargs are passed on to the generator function.
     """
     dist_shape = kwargs.pop("dist_shape", ())
-    one_d = _is_one_d(dist_shape)
     size = kwargs.pop("size", None)
     broadcast_shape = kwargs.pop("broadcast_shape", None)
     not_broadcast_kwargs = kwargs.pop("not_broadcast_kwargs", None)
@@ -1120,21 +1115,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/mixture.py

@@ -870,12 +870,6 @@ def random(self, point=None, size=None):
 
         # The `samples` array still has the `mixture_axis`, so we must remove it:
         output = samples[(..., 0) + (slice(None),) * len(event_shape)]
-
-        # Final oddity: if size == 1, pymc3 defaults to reducing the sample_shape dimension
-        # We do this to stay consistent with the rest of the package even though
-        # we shouldn't have to do it.
-        if size == 1:
-            output = output[0]
         return output
 
     def _distr_parameters_for_repr(self):

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):
@@ -1577,7 +1571,7 @@ def test_Triangular(
         assert prior["target"].shape == (prior_samples,) + shape
 
 
-def generate_shapes(include_params=False, xfail=False):
+def generate_shapes(include_params=False):
     # fmt: off
     mudim_as_event = [
         [None, 1, 3, 10, (10, 3), 100],
@@ -1596,20 +1590,13 @@ def generate_shapes(include_params=False, xfail=False):
         del mudim_as_event[-1]
         del mudim_as_dist[-1]
     data = itertools.chain(itertools.product(*mudim_as_event), itertools.product(*mudim_as_dist))
-    if xfail:
-        data = list(data)
-        for index in range(len(data)):
-            if data[index][0] in (None, 1):
-                data[index] = pytest.param(
-                    *data[index], marks=pytest.mark.xfail(reason="wait for PR #4214")
-                )
     return data
 
 
 class TestMvNormal(SeededTest):
     @pytest.mark.parametrize(
         ["sample_shape", "dist_shape", "mu_shape", "param"],
-        generate_shapes(include_params=True, xfail=False),
+        generate_shapes(include_params=True),
         ids=str,
     )
     def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param):
@@ -1619,7 +1606,7 @@ def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param):
 
     @pytest.mark.parametrize(
         ["sample_shape", "dist_shape", "mu_shape"],
-        generate_shapes(include_params=False, xfail=True),
+        generate_shapes(include_params=False),
         ids=str,
     )
     def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape):
@@ -1636,7 +1623,7 @@ def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape):
 
     @pytest.mark.parametrize(
         ["sample_shape", "dist_shape", "mu_shape"],
-        generate_shapes(include_params=False, xfail=True),
+        generate_shapes(include_params=False),
         ids=str,
     )
     def test_with_cov_rv(self, sample_shape, dist_shape, mu_shape):

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: