Infer steps from shape in GaussianRandomWalk

Author: canyon289

pymc/distributions/timeseries.py

@@ -18,6 +18,7 @@
 import numpy as np
 
 from aesara import scan
+from aesara.raise_op import Assert
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.utils import normalize_size_param
 
@@ -167,8 +168,26 @@ def dist(
 
         mu = at.as_tensor_variable(floatX(mu))
         sigma = at.as_tensor_variable(floatX(sigma))
+
+        # Check if shape contains information about number of steps
+        steps_from_shape = None
+        shape = kwargs.get("shape", None)
+        if shape is not None:
+            shape = to_tuple(shape)
+            if shape[-1] is not ...:
+                steps_from_shape = shape[-1] - 1
+
         if steps is None:
-            raise ValueError("Must specify steps parameter")
+            if steps_from_shape is not None:
+                steps = steps_from_shape
+            else:
+                raise ValueError("Must specify steps or shape parameter")
+        elif steps_from_shape is not None:
+            # Assert that steps and shape are consistent
+            steps = Assert(msg="Steps do not match last shape dimension")(
+                steps, at.eq(steps, steps_from_shape)
+            )
+
         steps = at.as_tensor_variable(intX(steps))
 
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma

pymc/tests/test_distributions_timeseries.py

@@ -86,10 +86,10 @@ def test_gaussian_random_walk_init_dist_shape(self, init):
         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)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=2)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=(5, 1))
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, shape=(5, 2))
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (5,)
 
         grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init=init)
@@ -113,6 +113,21 @@ def test_gaussianrandomwalk_broadcasted_by_init_dist(self):
         assert tuple(grw.shape.eval()) == (2, 3, 5)
         assert grw.eval().shape == (2, 3, 5)
 
+    @pytest.mark.parametrize("shape", ((6,), (3, 6)))
+    def test_inferred_steps_from_shape(self, shape):
+        x = GaussianRandomWalk.dist(shape=shape)
+        steps = x.owner.inputs[-1]
+        assert steps.eval() == 5
+
+    @pytest.mark.parametrize("shape", (None, (5, ...)))
+    def test_missing_steps(self, shape):
+        with pytest.raises(ValueError, match="Must specify steps or shape parameter"):
+            GaussianRandomWalk.dist(shape=shape)
+
+    def test_inconsistent_steps_and_shape(self):
+        with pytest.raises(AssertionError, match="Steps do not match last shape dimension"):
+            x = GaussianRandomWalk.dist(steps=12, shape=45)
+
     @pytest.mark.parametrize(
         "init",
         [