Add moment to GaussianRandomWalk and fix mu/sigma broadcasting bug

Author: ricardoV94

Diff for: pymc/distributions/timeseries.py

@@ -26,6 +26,7 @@
 from pymc.distributions import distribution, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
+from pymc.distributions.distribution import moment
 from pymc.distributions.logprob import ignore_logprob, logp
 from pymc.distributions.shape_utils import rv_size_is_none, to_tuple
 from pymc.util import check_dist_not_registered
@@ -131,7 +132,9 @@ def rng_fn(
         else:
             dist_shape = (*size, int(steps))
 
-        innovations = rng.normal(loc=mu, scale=sigma, size=dist_shape)
+        # Add one dimension to the right, so that mu and sigma broadcast safely along
+        # the steps dimension
+        innovations = rng.normal(loc=mu[..., None], scale=sigma[..., None], size=dist_shape)
         grw = np.concatenate([init[..., None], innovations], axis=-1)
         return np.cumsum(grw, axis=-1)
 
@@ -211,6 +214,14 @@ def dist(
 
         return super().dist([mu, sigma, init, steps], size=size, **kwargs)
 
+    def moment(rv, size, mu, sigma, init, steps):
+        grw_moment = at.zeros_like(rv)
+        grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init))
+        # Add one dimension to the right, so that mu broadcasts safely along the steps
+        # dimension
+        grw_moment = at.set_subtensor(grw_moment[..., 1:], mu[..., None])
+        return at.cumsum(grw_moment, axis=-1)
+
     def logp(
         value: at.Variable,
         mu: at.Variable,
@@ -225,7 +236,9 @@ def logp(
 
         # Make time series stationary around the mean value
         stationary_series = value[..., 1:] - value[..., :-1]
-        series_logp = logp(Normal.dist(mu, sigma), stationary_series)
+        # Add one dimension to the right, so that mu and sigma broadcast safely along
+        # the steps dimension
+        series_logp = logp(Normal.dist(mu[..., None], sigma[..., None]), stationary_series)
 
         return check_parameters(
             init_logp + series_logp.sum(axis=-1),

Diff for: pymc/tests/test_distributions_moments.py

@@ -103,7 +103,6 @@ def test_all_distributions_have_moments():
         dist_module.timeseries.AR,
         dist_module.timeseries.AR1,
         dist_module.timeseries.GARCH11,
-        dist_module.timeseries.GaussianRandomWalk,
         dist_module.timeseries.MvGaussianRandomWalk,
         dist_module.timeseries.MvStudentTRandomWalk,
     }

Diff for: pymc/tests/test_distributions_timeseries.py

@@ -29,6 +29,7 @@
 from pymc.model import Model
 from pymc.sampling import sample, sample_posterior_predictive
 from pymc.tests.helpers import select_by_precision
+from pymc.tests.test_distributions_moments import assert_moment_is_expected
 from pymc.tests.test_distributions_random import BaseTestDistributionRandom
 
 
@@ -142,6 +143,28 @@ def test_gaussian_random_walk_init_dist_logp(self, init):
             pm.logp(init, 0).eval() + scipy.stats.norm.logpdf(0),
         )
 
+    @pytest.mark.parametrize(
+        "mu, sigma, init, steps, size, expected",
+        [
+            (0, 1, Normal.dist(1), 10, None, np.ones((11,))),
+            (1, 1, Normal.dist(0), 10, (2,), np.full((2, 11), np.arange(11))),
+            (1, 1, Normal.dist([0, 1]), 10, None, np.vstack((np.arange(11), np.arange(11) + 1))),
+            (0, [1, 1], Normal.dist(0), 10, None, np.zeros((2, 11))),
+            (
+                [1, -1],
+                1,
+                Normal.dist(0),
+                10,
+                (4, 2),
+                np.full((4, 2, 11), np.vstack((np.arange(11), -np.arange(11)))),
+            ),
+        ],
+    )
+    def test_moment(self, mu, sigma, init, steps, size, expected):
+        with Model() as model:
+            GaussianRandomWalk("x", mu=mu, sigma=sigma, init=init, steps=steps, size=size)
+        assert_moment_is_expected(model, expected)
+
 
 @pytest.mark.xfail(reason="Timeseries not refactored")
 def test_AR():