Fix init dist and add tests

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -104,19 +104,20 @@ def rng_fn(
 
         # If size is None then the returned series should be (1+steps,)
         if size is None:
-            init_size = 1
-            steps_size = steps
+            bcast_shape = np.broadcast_shapes(
+                np.asarray(mu).shape,
+                np.asarray(sigma).shape,
+                np.asarray(init).shape,
+            )
+            dist_shape = (*bcast_shape, int(steps))
 
         # If size is None then the returned series should be (size, 1+steps)
         else:
             init_size = (*size, 1)
-            steps_size = (*size, steps)
-
-        init = np.reshape(init, init_size)
-        steps = rng.normal(loc=mu, scale=sigma, size=steps_size)
-
-        grw = np.concatenate([init, steps], axis=-1)
+            dist_shape = (*size, int(steps))
 
+        innovations = rng.normal(loc=mu, scale=sigma, size=dist_shape)
+        grw = np.concatenate([init[..., None], innovations], axis=-1)
         return np.cumsum(grw, axis=-1)
 
 
@@ -149,7 +150,8 @@ class GaussianRandomWalk(distribution.Continuous):
     rv_op = gaussianrandomwalk
 
     def __new__(cls, name, mu=0.0, sigma=1.0, init=None, steps=None, **kwargs):
-        check_dist_not_registered(init)
+        if init is not None:
+            check_dist_not_registered(init)
         return super().__new__(cls, name, mu, sigma, init, steps, **kwargs)
 
     @classmethod
@@ -163,14 +165,15 @@ def dist(
             raise ValueError("Must specify steps parameter")
         steps = at.as_tensor_variable(intX(steps))
 
-        if "shape" in kwargs.keys():
-            shape = kwargs["shape"]
+        shape = kwargs.get("shape", None)
+        if size is None and shape is None:
+            init_size = None
         else:
-            shape = None
+            init_size = to_tuple(size) if size is not None else to_tuple(shape)[:-1]
 
-        # If no scalar distribution is passed then initialize with a Normal of same sd and mu
+        # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
         if init is None:
-            init = Normal.dist(mu, sigma, size=size)
+            init = Normal.dist(mu, sigma, size=init_size)
         else:
             if not (
                 isinstance(init, at.TensorVariable)
@@ -180,12 +183,12 @@ def dist(
             ):
                 raise TypeError("init must be a univariate distribution variable")
 
-            if size is not None or shape is not None:
-                init = change_rv_size(init, to_tuple(size or shape))
+            if init_size is not None:
+                init = change_rv_size(init, init_size)
             else:
-                # If not explicit, size is determined by the shape of mu and sigma
-                mu_ = at.broadcast_arrays(mu, sigma)[0]
-                init = change_rv_size(init, mu_.shape)
+                # If not explicit, size is determined by the shapes of mu, sigma, and init
+                bcast_shape = at.broadcast_arrays(mu, sigma, init)[0].shape
+                init = change_rv_size(init, bcast_shape)
 
         # Ignores logprob of init var because that's accounted for in the logp method
         init.tag.ignore_logprob = True

pymc/tests/test_distributions_timeseries.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 import numpy as np
 import pytest
+import scipy.stats
 
 import pymc as pm
 
@@ -77,6 +78,52 @@ def test_gaussianrandomwalk_inference():
     np.testing.assert_allclose([mu, sigma], [recovered_mu, recovered_sigma], atol=0.2)
 
 
+@pytest.mark.parametrize("init", [None, pm.Normal.dist()])
+def test_gaussian_random_walk_init_dist_shape(init):
+    """Test that init_dist is properly resized"""
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
+
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, size=(5,))
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,)
+
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=1)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
+
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=(5, 1))
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,)
+
+    grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init=init)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,)
+
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init=init)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,)
+
+    grw = pm.GaussianRandomWalk.dist(mu=np.zeros((3, 1)), sigma=[1, 1], steps=1, init=init)
+    assert tuple(grw.owner.inputs[-2].shape.eval()) == (3, 2)
+
+
+def test_gaussianrandomwalk_broadcasted_by_init_dist():
+    grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=4, init=pm.Normal.dist(size=(2, 3)))
+    assert tuple(grw.shape.eval()) == (2, 3, 5)
+    assert grw.eval().shape == (2, 3, 5)
+
+
+@pytest.mark.parametrize(
+    "init",
+    [
+        pm.HalfNormal.dist(sigma=2),
+        pm.StudentT.dist(nu=4, mu=1, sigma=0.5),
+    ],
+)
+def test_gaussian_random_walk_init_dist_logp(init):
+    grw = pm.GaussianRandomWalk.dist(init=init, steps=1)
+    assert np.isclose(
+        pm.logp(grw, [0, 0]).eval(),
+        pm.logp(init, 0).eval() + scipy.stats.norm.logpdf(0),
+    )
+
+
 @pytest.mark.xfail(reason="Timeseries not refactored")
 def test_AR():
     # AR1