Apply suggestions from code review

Co-authored-by: Ricardo Vieira <[email protected]>

Author: junpenglao

pymc/distributions/timeseries.py

@@ -690,7 +690,7 @@ def rv_op(cls, omega, alpha_1, beta_1, initial_vol, init_dist, steps, size=None)
         init_dist = change_dist_size(init_dist, batch_size)
         # initial_vol = initial_vol * at.ones(batch_size)
 
-        # Create OpFromGraph representing random draws form AR process
+        # Create OpFromGraph representing random draws from GARCH11 process
         # Variables with underscore suffix are dummy inputs into the OpFromGraph
         init_ = init_dist.type()
         initial_vol_ = initial_vol.type()
@@ -701,8 +701,7 @@ def rv_op(cls, omega, alpha_1, beta_1, initial_vol, init_dist, steps, size=None)
 
         noise_rng = aesara.shared(np.random.default_rng())
 
-        def step(*args):
-            prev_y, prev_sigma, omega, alpha_1, beta_1, rng = args
+        def step(prev_y, prev_sigma, omega, alpha_1, beta_1, rng):
             new_sigma = at.sqrt(
                 omega + alpha_1 * at.square(prev_y) + beta_1 * at.square(prev_sigma)
             )
@@ -761,6 +760,7 @@ def volatility_update(x, vol, w, a, b):
         sequences=[value_dimswapped[:-1]],
         outputs_info=[initial_vol],
         non_sequences=[omega, alpha_1, beta_1],
+       strict = True,
     )
     sigma_t = at.concatenate([[initial_vol], vol])
     # Compute and collapse logp across time dimension

pymc/tests/distributions/test_timeseries.py

@@ -533,10 +533,11 @@ def test_logp(self):
         np.testing.assert_allclose(garch_like, reg_like, 10 ** (-decimal))
 
     @pytest.mark.parametrize(
-        "arg_name",
+        "batched_param",
         ["omega", "alpha_1", "beta_1", "initial_vol"],
     )
-    def test_batched_size(self, arg_name):
+    @pytest.mark.parametrize("explicit_shape", (True, False))
+    def test_batched_size(self, explicit_shape, batched_param):
         steps, batch_size = 100, 5
         param_val = np.square(np.random.randn(batch_size))
         init_kwargs = dict(