Repository: pymc-devs/pymc

Issue:
Improve default GaussianRandomWalk and AR init distributions
When refactoring the GaussianRandomWalk to V4, we set the `init` to default to `Normal(mu, sigma)`, which does not make much sense. `mu` is the drift parameter, and says nothing about the expected starting location of the series. `sigma` probably also makes little sense, as it assumes the uncertainty about the initial location has anything to do with the innovations noise. 

Possible alternatives:
* `init = Normal.dist(0, 1)`
* `init = Flat.dist()`, which was the default in V3, but it means that prior and posterior predictive will fail by default.

We should also rename `init` to `init_dist` to emphasize this parameter must de a distribution.

https://github.com/pymc-devs/pymc/blob/c76b9b9e74ee6f13993f255dd09d9290b882d6ab/pymc/distributions/timeseries.py#L239

A similar question arises in the AR, where we again reuse the innovation noise `sigma` for the uncertainty about the initial location of the series.

https://github.com/pymc-devs/pymc/blob/c76b9b9e74ee6f13993f255dd09d9290b882d6ab/pymc/distributions/timeseries.py#L429

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5829?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5829](https://codecov.io/gh/pymc-devs/pymc/pull/5829?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (193147c) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/6b22ed51599f4d41328e1b28d454f09ac21bc877?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (6b22ed5) will **decrease** coverage by `15.30%`.
> The diff coverage is `10.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5829/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5829       +/-   ##
===========================================
- Coverage   89.40%   74.10%   -15.31%     
===========================================
  Files          74       74               
  Lines       13772    13777        +5     
===========================================
- Hits        12313    10209     -2104     
- Misses       1459     3568     +2109     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/distributions/timeseries.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL3RpbWVzZXJpZXMucHk=) | `30.10% <10.00%> (-47.37%)` | :arrow_down: |
| [pymc/variational/updates.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC91cGRhdGVzLnB5) | `11.33% <0.00%> (-80.79%)` | :arrow_down: |
| [pymc/model\_graph.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9tb2RlbF9ncmFwaC5weQ==) | `17.60% <0.00%> (-76.00%)` | :arrow_down: |
| [pymc/variational/inference.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC9pbmZlcmVuY2UucHk=) | `19.38% <0.00%> (-67.86%)` | :arrow_down: |
| [pymc/printing.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9wcmludGluZy5weQ==) | `21.78% <0.00%> (-64.36%)` | :arrow_down: |
| [pymc/variational/test\_functions.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC90ZXN0X2Z1bmN0aW9ucy5weQ==) | `38.09% <0.00%> (-61.91%)` | :arrow_down: |
| [pymc/distributions/bound.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2JvdW5kLnB5) | `45.54% <0.00%> (-54.46%)` | :arrow_down: |
| [pymc/variational/stein.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC9zdGVpbi5weQ==) | `48.33% <0.00%> (-51.67%)` | :arrow_down: |
| [pymc/distributions/censored.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2NlbnNvcmVkLnB5) | `44.18% <0.00%> (-48.84%)` | :arrow_down: |
| [pymc/variational/opvi.py](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC9vcHZpLnB5) | `34.41% <0.00%> (-48.04%)` | :arrow_down: |
| ... and [28 more](https://codecov.io/gh/pymc-devs/pymc/pull/5829/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |


Patch:
diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py
index 2a36e1f39..6aa902bcc 100644
--- a/pymc/distributions/timeseries.py
+++ b/pymc/distributions/timeseries.py
@@ -125,22 +125,24 @@ class GaussianRandomWalkRV(RandomVariable):
     dtype = "floatX"
     _print_name = ("GaussianRandomWalk", "\\operatorname{GaussianRandomWalk}")
 
-    def make_node(self, rng, size, dtype, mu, sigma, init, steps):
+    def make_node(self, rng, size, dtype, mu, sigma, init_dist, steps):
         steps = at.as_tensor_variable(steps)
         if not steps.ndim == 0 or not steps.dtype.startswith("int"):
             raise ValueError("steps must be an integer scalar (ndim=0).")
 
         mu = at.as_tensor_variable(mu)
         sigma = at.as_tensor_variable(sigma)
-        init = at.as_tensor_variable(init)
+        init_dist = at.as_tensor_variable(init_dist)
 
         # Resize init distribution
         size = normalize_size_param(size)
         # If not explicit, size is determined by the shapes of mu, sigma, and init
-        init_size = size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init)
-        init = change_rv_size(init, init_size)
+        init_dist_size = (
+            size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init_dist)
+        )
+        init_dist = change_rv_size(init_dist, init_dist_size)
 
-        return super().make_node(rng, size, dtype, mu, sigma, init, steps)
+        return super().make_node(rng, size, dtype, mu, sigma, init_dist, steps)
 
     def _supp_shape_from_params(self, dist_params, reop_param_idx=0, param_shapes=None):
         steps = dist_params[3]
@@ -153,7 +155,7 @@ class GaussianRandomWalkRV(RandomVariable):
         rng: np.random.RandomState,
         mu: Union[np.ndarray, float],
         sigma: Union[np.ndarray, float],
-        init: float,
+        init_dist: Union[np.ndarray, float],
         steps: int,
         size: Tuple[int],
     ) -> np.ndarray:
@@ -170,16 +172,16 @@ class GaussianRandomWalkRV(RandomVariable):
         ----------
         rng: np.random.RandomState
            Numpy random number generator
-        mu: array_like
+        mu: array_like of float
            Random walk mean
-        sigma: np.ndarray
+        sigma: array_like of float
             Standard deviation of innovation (sigma > 0)
-        init: float
+        init_dist: array_like of float
             Initialization value for GaussianRandomWalk
         steps: int
             Length of random walk, must be greater than 1. Returned array will be of size+1 to
             account as first value is initial value
-        size: int
+        size: tuple of int
             The number of Random Walk time series generated
 
         Returns
@@ -196,7 +198,7 @@ class GaussianRandomWalkRV(RandomVariable):
             bcast_shape = np.broadcast_shapes(
                 np.asarray(mu).shape,
                 np.asarray(sigma).shape,
-                np.asarray(init).shape,
+                np.asarray(init_dist).shape,
             )
             dist_shape = (*bcast_shape, int(steps))
 
@@ -207,7 +209,7 @@ class GaussianRandomWalkRV(RandomVariable):
         # 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)
+        grw = np.concatenate([init_dist[..., None], innovations], axis=-1)
         return np.cumsum(grw, axis=-1)
 
 
@@ -223,15 +225,14 @@ class GaussianRandomWalk(distribution.Continuous):
         innovation drift, defaults to 0.0
     sigma : tensor_like of float, optional
         sigma > 0, innovation standard deviation, defaults to 1.0
-    init : unnamed distribution
+    init_dist : unnamed distribution
         Univariate distribution of the initial value, created with the `.dist()` API.
-        Defaults to a unit Normal.
 
         .. 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
@@ -247,43 +248,53 @@ class GaussianRandomWalk(distribution.Continuous):
         return super().__new__(cls, *args, steps=steps, **kwargs)
 
     @classmethod
-    def dist(
-        cls, mu=0.0, sigma=1.0, *, init=None, steps=None, size=None, **kwargs
-    ) -> at.TensorVariable:
+    def dist(cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, **kwargs) -> at.TensorVariable:
 
         mu = at.as_tensor_variable(floatX(mu))
         sigma = at.as_tensor_variable(floatX(sigma))
 
         steps = get_steps(
             steps=steps,
-            shape=kwargs.get("shape", None),
+            shape=kwargs.get("shape"),
             step_shape_offset=1,
         )
         if steps is None:
             raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps))
 
+        if "init" in kwargs:
+            warnings.warn(
+                "init parameter is now called init_dist. Using init will raise an error in a future release.",
+                FutureWarning,
+            )
+            init_dist = kwargs.pop("init")
+
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
-        if init is None:
-            init = Normal.dist(0, 1)
+        if init_dist is None:
+            warnings.warn(
+                "Initial distribution not specified, defaulting to `Normal.dist(0, 100)`."
+                "You can specify an init_dist manually to suppress this warning.",
+                UserWarning,
+            )
+            init_dist = Normal.dist(0, 100)
         else:
             if not (
-                isinstance(init, at.TensorVariable)
-                and init.owner is not None
-                and isinstance(init.owner.op, RandomVariable)
-                and init.owner.op.ndim_supp == 0
+                isinstance(init_dist, at.TensorVariable)
+                and init_dist.owner is not None
+                and isinstance(init_dist.owner.op, RandomVariable)
+                and init_dist.owner.op.ndim_supp == 0
             ):
                 raise TypeError("init must be a univariate distribution variable")
-            check_dist_not_registered(init)
+            check_dist_not_registered(init_dist)
 
         # Ignores logprob of init var because that's accounted for in the logp method
-        init = ignore_logprob(init)
+        init_dist = ignore_logprob(init_dist)
 
-        return super().dist([mu, sigma, init, steps], size=size, **kwargs)
+        return super().dist([mu, sigma, init_dist, steps], **kwargs)
 
-    def moment(rv, size, mu, sigma, init, steps):
+    def moment(rv, size, mu, sigma, init_dist, steps):
         grw_moment = at.zeros_like(rv)
-        grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init))
+        grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init_dist))
         # 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])
@@ -293,13 +304,13 @@ class GaussianRandomWalk(distribution.Continuous):
         value: at.Variable,
         mu: at.Variable,
         sigma: at.Variable,
-        init: at.Variable,
+        init_dist: at.Variable,
         steps: at.Variable,
     ) -> at.TensorVariable:
         """Calculate log-probability of Gaussian Random Walk distribution at specified value."""
 
         # Calculate initialization logp
-        init_logp = logp(init, value[..., 0])
+        init_logp = logp(init_dist, value[..., 0])
 
         # Make time series stationary around the mean value
         stationary_series = value[..., 1:] - value[..., :-1]
@@ -360,10 +371,10 @@ class AR(SymbolicDistribution):
     constant: bool, optional
         Whether the first element of rho should be used as a constant term in the AR
         process. Defaults to False
-    init_dist: unnamed distribution, optional
-        Scalar or vector distribution for initial values. Defaults to a unit Normal.
-        Distribution should be  created via the `.dist()` API, and have  dimension
-        (*size, ar_order). If not, it will be automatically resized.
+    init_dist: unnamed distribution
+        Scalar or vector distribution for initial values. Distribution should be
+        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.
 
@@ -371,8 +382,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
     -----
@@ -429,7 +439,7 @@ class AR(SymbolicDistribution):
                 "init parameter is now called init_dist. Using init will raise an error in a future release.",
                 FutureWarning,
             )
-            init_dist = kwargs["init"]
+            init_dist = kwargs.pop("init")
 
         ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
         steps = get_steps(steps=steps, shape=kwargs.get("shape", None), step_shape_offset=ar_order)
@@ -452,7 +462,13 @@ class AR(SymbolicDistribution):
                     f"got ndim_supp={init_dist.owner.op.ndim_supp}.",
                 )
         else:
-            init_dist = Normal.dist(0, 1, size=(*sigma.shape, ar_order))
+            warnings.warn(
+                "Initial distribution not specified, defaulting to "
+                "`Normal.dist(0, 100, shape=...)`. You can specify an init_dist "
+                "manually to suppress this warning.",
+                UserWarning,
+            )
+            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)


Test patch:
diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py
index c805e01f1..ba1611404 100644
--- a/pymc/tests/test_distributions.py
+++ b/pymc/tests/test_distributions.py
@@ -2610,7 +2610,7 @@ class TestMatchesScipy:
         def ref_logp(value, mu, sigma, steps):
             # Relying on fact that init will be normal by default
             return (
-                scipy.stats.norm.logpdf(value[0])
+                scipy.stats.norm.logpdf(value[0], 0, 100)  # default init_dist has a scale 100
                 + scipy.stats.norm.logpdf(np.diff(value), mu, sigma).sum()
             )
 
diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py
index 9be401bcf..00fbecae0 100644
--- a/pymc/tests/test_distributions_timeseries.py
+++ b/pymc/tests/test_distributions_timeseries.py
@@ -100,8 +100,13 @@ class TestGaussianRandomWalk:
         size = None
 
         pymc_dist = pm.GaussianRandomWalk
-        pymc_dist_params = {"mu": 1.0, "sigma": 2, "init": pm.Constant.dist(0), "steps": 4}
-        expected_rv_op_params = {"mu": 1.0, "sigma": 2, "init": pm.Constant.dist(0), "steps": 4}
+        pymc_dist_params = {"mu": 1.0, "sigma": 2, "init_dist": pm.Constant.dist(0), "steps": 4}
+        expected_rv_op_params = {
+            "mu": 1.0,
+            "sigma": 2,
+            "init_dist": pm.Constant.dist(0),
+            "steps": 4,
+        }
 
         checks_to_run = [
             "check_pymc_params_match_rv_op",
@@ -142,36 +147,38 @@ class TestGaussianRandomWalk:
     @pytest.mark.parametrize("init", [None, pm.Normal.dist()])
     def test_gaussian_random_walk_init_dist_shape(self, init):
         """Test that init_dist is properly resized"""
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == ()
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init=init, size=(5,))
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=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=2)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=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, 2))
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=1, init_dist=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)
+        grw = pm.GaussianRandomWalk.dist(mu=[0, 0], sigma=1, steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (2,)
 
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init=init)
+        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=[1, 1], steps=1, init_dist=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)
+        grw = pm.GaussianRandomWalk.dist(mu=np.zeros((3, 1)), sigma=[1, 1], steps=1, init_dist=init)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (3, 2)
 
     def test_shape_ellipsis(self):
         grw = pm.GaussianRandomWalk.dist(
-            mu=0, sigma=1, steps=5, init=pm.Normal.dist(), shape=(3, ...)
+            mu=0, sigma=1, steps=5, init_dist=pm.Normal.dist(), shape=(3, ...)
         )
         assert tuple(grw.shape.eval()) == (3, 6)
         assert tuple(grw.owner.inputs[-2].shape.eval()) == (3,)
 
     def test_gaussianrandomwalk_broadcasted_by_init_dist(self):
-        grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=4, init=pm.Normal.dist(size=(2, 3)))
+        grw = pm.GaussianRandomWalk.dist(
+            mu=0, sigma=1, steps=4, init_dist=pm.Normal.dist(size=(2, 3))
+        )
         assert tuple(grw.shape.eval()) == (2, 3, 5)
         assert grw.eval().shape == (2, 3, 5)
 
@@ -210,14 +217,14 @@ class TestGaussianRandomWalk:
         ],
     )
     def test_gaussian_random_walk_init_dist_logp(self, init):
-        grw = pm.GaussianRandomWalk.dist(init=init, steps=1)
+        grw = pm.GaussianRandomWalk.dist(init_dist=init, steps=1)
         assert np.isclose(
             pm.logp(grw, [0, 0]).eval(),
             pm.logp(init, 0).eval() + scipy.stats.norm.logpdf(0),
         )
 
     @pytest.mark.parametrize(
-        "mu, sigma, init, steps, size, expected",
+        "mu, sigma, init_dist, 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))),
@@ -233,11 +240,15 @@ class TestGaussianRandomWalk:
             ),
         ],
     )
-    def test_moment(self, mu, sigma, init, steps, size, expected):
+    def test_moment(self, mu, sigma, init_dist, steps, size, expected):
         with Model() as model:
-            GaussianRandomWalk("x", mu=mu, sigma=sigma, init=init, steps=steps, size=size)
+            GaussianRandomWalk("x", mu=mu, sigma=sigma, init_dist=init_dist, steps=steps, size=size)
         assert_moment_is_expected(model, expected)
 
+    def test_init_deprecated_arg(self):
+        with pytest.warns(FutureWarning, match="init parameter is now called init_dist"):
+            pm.GaussianRandomWalk.dist(init=Normal.dist(), shape=(10,))
+
 
 class TestAR:
     def test_order1_logp(self):
@@ -304,11 +315,25 @@ class TestAR:
         y_tp = np.random.randn(batch_size, steps)
         with Model() as t0:
             beta = Normal("beta", 0.0, 1.0, shape=(batch_size, ar_order), initval=beta_tp)
-            AR("y", beta, sigma=1.0, shape=(batch_size, steps), initval=y_tp)
+            AR(
+                "y",
+                beta,
+                sigma=1.0,
+                init_dist=Normal.dist(0, 1),
+                shape=(batch_size, steps),
+                initval=y_tp,
+            )
         with Model() as t1:
             beta = Normal("beta", 0.0, 1.0, shape=(batch_size, ar_order), initval=beta_tp)
             for i in range(batch_size):
-                AR(f"y_{i}", beta[i], sigma=1.0, shape=steps, initval=y_tp[i])
+                AR(
+                    f"y_{i}",
+                    beta[i],
+                    init_dist=Normal.dist(0, 1),
+                    sigma=1.0,
+                    shape=steps,
+                    initval=y_tp[i],
+                )
 
         np.testing.assert_allclose(
             t0.compile_logp()(t0.initial_point()),
@@ -373,7 +398,7 @@ class TestAR:
         beta_tp = aesara.shared(np.random.randn(ar_order), shape=(3,))
         y_tp = np.random.randn(batch_size, steps)
         with Model() as t0:
-            init_dist = Normal.dist(0.0, 1.0, size=(batch_size, ar_order))
+            init_dist = Normal.dist(0.0, 100.0, size=(batch_size, ar_order))
             AR("y", beta_tp, sigma=0.01, init_dist=init_dist, steps=steps, initval=y_tp)
         with Model() as t1:
             for i in range(batch_size):
@@ -434,6 +459,10 @@ class TestAR:
             AR("x", rho=[0, 0], init_dist=init_dist, steps=5, size=size)
         assert_moment_is_expected(model, expected, check_finite_logp=False)
 
+    def test_init_deprecated_arg(self):
+        with pytest.warns(FutureWarning, match="init parameter is now called init_dist"):
+            pm.AR.dist(rho=[1, 2, 3], init=Normal.dist(), shape=(10,))
+
 
 @pytest.mark.xfail(reason="Timeseries not refactored", raises=NotImplementedError)
 def test_GARCH11():