From ca1a74d8a54fbb0147d9066029bbb949217b489f Mon Sep 17 00:00:00 2001
From: Ricardo <ricardo.vieira1994@gmail.com>
Date: Thu, 2 Jun 2022 11:05:28 +0200
Subject: [PATCH] Remove user-facing mentions of size in timeseries
 distributions

---
 pymc/distributions/timeseries.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py
index 194186bb7c2..6fc116a64ca 100644
--- a/pymc/distributions/timeseries.py
+++ b/pymc/distributions/timeseries.py
@@ -231,8 +231,8 @@ class GaussianRandomWalk(distribution.Continuous):
         .. warning:: init will be cloned, rendering them independent of the ones passed as input.
 
     steps : int, optional
-        Number of steps in Gaussian Random Walk (steps > 0). Only needed if size is
-        used to specify distribution
+        Number of steps in Gaussian Random Walk (steps > 0). Only needed if shape is not
+        provided.
     """
 
     rv_op = gaussianrandomwalk
@@ -249,7 +249,7 @@ def __new__(cls, *args, steps=None, **kwargs):
 
     @classmethod
     def dist(
-        cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, size=None, **kwargs
+        cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, **kwargs
     ) -> at.TensorVariable:
 
         mu = at.as_tensor_variable(floatX(mu))
@@ -257,7 +257,7 @@ def dist(
 
         steps = get_steps(
             steps=steps,
-            shape=kwargs.get("shape", None),
+            shape=kwargs.get("shape"),
             step_shape_offset=1,
         )
         if steps is None:
@@ -292,7 +292,7 @@ def dist(
         # Ignores logprob of init var because that's accounted for in the logp method
         init_dist = ignore_logprob(init_dist)
 
-        return super().dist([mu, sigma, init_dist, steps], size=size, **kwargs)
+        return super().dist([mu, sigma, init_dist, steps], **kwargs)
 
     def moment(rv, size, mu, sigma, init_dist, steps):
         grw_moment = at.zeros_like(rv)
@@ -375,7 +375,7 @@ class AR(SymbolicDistribution):
         process. Defaults to False
     init_dist: unnamed distribution
         Scalar or vector distribution for initial values. Distribution should be
-        created via the `.dist()` API, and have  dimension (*size, ar_order). If not,
+        created via the `.dist()` API, and have  shape (*shape[:-1], ar_order). If not,
         it will be automatically resized.
 
         .. warning:: init_dist will be cloned, rendering it independent of the one passed as input.
@@ -384,8 +384,7 @@ class AR(SymbolicDistribution):
         Order of the AR process. Inferred from length of the last dimension of rho, if
         possible. ar_order = rho.shape[-1] if constant else rho.shape[-1] - 1
     steps : int, optional
-        Number of steps in AR process (steps > 0). Only needed if size is used to
-        specify distribution
+        Number of steps in AR process (steps > 0). Only needed if shape is not provided.
 
     Notes
     -----
@@ -471,7 +470,7 @@ def dist(
                 "manually to suppress this warning.",
                 UserWarning,
             )
-            init_dist = Normal.dist(0, 100, size=(*sigma.shape, ar_order))
+            init_dist = Normal.dist(0, 100, shape=(*sigma.shape, ar_order))
 
         # Tell Aeppl to ignore init_dist, as it will be accounted for in the logp term
         init_dist = ignore_logprob(init_dist)