Switch to new Scan API

Co-authored-by: Michal-Novomestsky

Author: ricardoV94

pymc/distributions/timeseries.py

@@ -442,23 +442,26 @@ def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None
         rhos_bcast = pt.broadcast_to(rhos, rhos_bcast_shape)
 
         def step(*args):
-            *prev_xs, reversed_rhos, sigma, rng = args
+            *prev_xs, rng, reversed_rhos, sigma = args
             if constant_term:
                 mu = reversed_rhos[-1] + pt.sum(prev_xs * reversed_rhos[:-1], axis=0)
             else:
                 mu = pt.sum(prev_xs * reversed_rhos, axis=0)
             next_rng, new_x = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
-            return new_x, {rng: next_rng}
+            return new_x, next_rng
 
         # We transpose inputs as scan iterates over first dimension
-        innov, innov_updates = pytensor.scan(
+        innov, noise_next_rng = pytensor.scan(
             fn=step,
-            outputs_info=[{"initial": init_dist.T, "taps": range(-ar_order, 0)}],
-            non_sequences=[rhos_bcast.T[::-1], sigma.T, noise_rng],
+            outputs_info=[
+                {"initial": init_dist.T, "taps": range(-ar_order, 0)},
+                noise_rng,
+            ],
+            non_sequences=[rhos_bcast.T[::-1], sigma.T],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
         ar = pt.concatenate([init_dist, innov.T], axis=-1)
 
         return AutoRegressiveRV(
@@ -710,24 +713,25 @@ def rv_op(cls, omega, alpha_1, beta_1, initial_vol, init_dist, steps, size=None)
 
         # Create OpFromGraph representing random draws from GARCH11 process
 
-        def step(prev_y, prev_sigma, omega, alpha_1, beta_1, rng):
+        def step(prev_y, prev_sigma, rng, omega, alpha_1, beta_1):
             new_sigma = pt.sqrt(
                 omega + alpha_1 * pt.square(prev_y) + beta_1 * pt.square(prev_sigma)
             )
             next_rng, new_y = Normal.dist(mu=0, sigma=new_sigma, rng=rng).owner.outputs
-            return (new_y, new_sigma), {rng: next_rng}
+            return new_y, new_sigma, next_rng
 
-        (y_t, _), innov_updates = pytensor.scan(
+        y_t, _, noise_next_rng = pytensor.scan(
             fn=step,
             outputs_info=[
                 init_dist,
                 pt.broadcast_to(initial_vol.astype("floatX"), init_dist.shape),
+                noise_rng,
             ],
-            non_sequences=[omega, alpha_1, beta_1, noise_rng],
+            non_sequences=[omega, alpha_1, beta_1],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
 
         garch11 = pt.concatenate([init_dist[None, ...], y_t], axis=0).dimshuffle(
             (*range(1, y_t.ndim), 0)
@@ -816,12 +820,13 @@ def garch11_logp(
     def volatility_update(x, vol, w, a, b):
         return pt.sqrt(w + a * pt.square(x) + b * pt.square(vol))
 
-    vol, _ = pytensor.scan(
+    vol = pytensor.scan(
         fn=volatility_update,
         sequences=[value_dimswapped[:-1]],
         outputs_info=[initial_vol],
         non_sequences=[omega, alpha_1, beta_1],
         strict=True,
+        return_updates=False,
     )
     sigma_t = pt.concatenate([[initial_vol], vol])
     # Compute and collapse logp across time dimension
@@ -861,21 +866,21 @@ def rv_op(cls, init_dist, steps, sde_pars, dt, sde_fn, size=None):
 
         # Create OpFromGraph representing random draws from SDE process
         def step(*prev_args):
-            prev_y, *prev_sde_pars, rng = prev_args
+            prev_y, rng, *prev_sde_pars = prev_args
             f, g = sde_fn(prev_y, *prev_sde_pars)
             mu = prev_y + dt * f
             sigma = pt.sqrt(dt) * g
             next_rng, next_y = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
-            return next_y, {rng: next_rng}
+            return next_y, next_rng
 
-        y_t, innov_updates = pytensor.scan(
+        y_t, noise_next_rng = pytensor.scan(
             fn=step,
-            outputs_info=[init_dist],
-            non_sequences=[*sde_pars, noise_rng],
+            outputs_info=[init_dist, noise_rng],
+            non_sequences=[*sde_pars],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
 
         sde_out = pt.concatenate([init_dist[None, ...], y_t], axis=0).dimshuffle(
             (*range(1, y_t.ndim), 0)

pymc/logprob/transforms.py

@@ -729,11 +729,12 @@ def calc_delta_x(value, prior_result):
                 2 * prior_result * pt.erfcx(prior_result) - 2 / pt.sqrt(np.pi)
             )
 
-        result, updates = scan(
+        result = scan(
             fn=calc_delta_x,
             outputs_info=pt.ones_like(x),
             non_sequences=value,
             n_steps=10,
+            return_updates=False,
         )
         return result[-1]
 

pymc/pytensorf.py

@@ -331,7 +331,7 @@ def jacobian1(f, v):
     def grad_i(i):
         return gradient1(f[i], v)
 
-    return pytensor.map(grad_i, idx)[0]
+    return pytensor.map(grad_i, idx, return_updates=False)
 
 
 def jacobian(f, vars=None):
@@ -355,8 +355,13 @@ def grad_ii(i, f, x):
         return grad(f[i], x)[i]
 
     return pytensor.scan(
-        grad_ii, sequences=[idx], n_steps=f.shape[0], non_sequences=[f, x], name="jacobian_diag"
-    )[0]
+        grad_ii,
+        sequences=[idx],
+        n_steps=f.shape[0],
+        non_sequences=[f, x],
+        name="jacobian_diag",
+        return_updates=False,
+    )
 
 
 @pytensor.config.change_flags(compute_test_value="ignore")
@@ -381,7 +386,7 @@ def hessian_diag1(f, v):
     def hess_ii(i):
         return gradient1(g[i], v)[i]
 
-    return pytensor.map(hess_ii, idx)[0]
+    return pytensor.map(hess_ii, idx, return_updates=False)
 
 
 @pytensor.config.change_flags(compute_test_value="ignore")

pymc/variational/approximations.py

@@ -392,7 +392,10 @@ def evaluate_over_trace(self, node):
         def sample(post, *_):
             return graph_replace(node, {self.input: post}, strict=False)
 
-        nodes, _ = pytensor.scan(
-            sample, self.histogram, non_sequences=_known_scan_ignored_inputs(makeiter(node))
+        nodes = pytensor.scan(
+            sample,
+            self.histogram,
+            non_sequences=_known_scan_ignored_inputs(makeiter(node)),
+            return_updates=False,
         )
         return nodes

pymc/variational/opvi.py

@@ -1062,8 +1062,11 @@ def symbolic_sample_over_posterior(self, node):
         def sample(post, *_):
             return graph_replace(node, {self.input: post}, strict=False)
 
-        nodes, _ = pytensor.scan(
-            sample, random, non_sequences=_known_scan_ignored_inputs(makeiter(random))
+        nodes = pytensor.scan(
+            sample,
+            random,
+            non_sequences=_known_scan_ignored_inputs(makeiter(random)),
+            return_updates=False,
         )
         assert self.input not in set(pytensor.graph.graph_inputs(makeiter(nodes)))
         return nodes
@@ -1451,8 +1454,11 @@ def symbolic_sample_over_posterior(self, node, more_replacements=None):
         def sample(*post):
             return graph_replace(node, dict(zip(self.inputs, post)), strict=False)
 
-        nodes, _ = pytensor.scan(
-            sample, self.symbolic_randoms, non_sequences=_known_scan_ignored_inputs(makeiter(node))
+        nodes = pytensor.scan(
+            sample,
+            self.symbolic_randoms,
+            non_sequences=_known_scan_ignored_inputs(makeiter(node)),
+            return_updates=False,
         )
         assert not (set(self.inputs) & set(pytensor.graph.graph_inputs(makeiter(nodes))))
         return nodes

tests/distributions/test_custom.py

@@ -362,32 +362,40 @@ def scan_step(left, right):
             return x, x_update
 
         def dist(size):
-            xs, updates = scan(
-                fn=scan_step,
-                sequences=[
-                    pt.as_tensor_variable(np.array([-4, -3])),
-                    pt.as_tensor_variable(np.array([-2, -1])),
-                ],
-                name="xs",
-            )
+            with pytest.warns(DeprecationWarning, match="Scan return signature will change"):
+                xs, updates = scan(
+                    fn=scan_step,
+                    sequences=[
+                        pt.as_tensor_variable(np.array([-4, -3])),
+                        pt.as_tensor_variable(np.array([-2, -1])),
+                    ],
+                    name="xs",
+                    # There's a bug in the ordering of outputs when there's a mapped `None` output
+                    # We have to stick with the deprecated API for now
+                    return_updates=True,
+                )
             return xs
 
         with Model() as model:
             CustomDist("x", dist=dist)
         assert_support_point_is_expected(model, np.array([-3, -2]))
 
     def test_custom_dist_default_support_point_scan_recurring(self):
-        def scan_step(xtm1):
-            x = Normal.dist(xtm1 + 1)
-            x_update = collect_default_updates([x])
-            return x, x_update
+        def scan_step(xtm1, rng):
+            next_rng, x = Normal.dist(xtm1 + 1, rng=rng).owner.outputs
+            return x, next_rng
 
         def dist(size):
-            xs, _ = scan(
+            rng = pytensor.shared(np.random.default_rng())
+            xs, _next_rng = scan(
                 fn=scan_step,
-                outputs_info=pt.as_tensor_variable(np.array([0])).astype(float),
+                outputs_info=[
+                    pt.as_tensor_variable(np.array([0])).astype(float),
+                    rng,
+                ],
                 n_steps=3,
                 name="xs",
+                return_updates=False,
             )
             return xs
 
@@ -527,16 +535,20 @@ def test_scan(self):
         def trw(nu, sigma, steps, size):
             if rv_size_is_none(size):
                 size = ()
+            rng = pytensor.shared(np.random.default_rng())
 
-            def step(xtm1, nu, sigma):
-                x = StudentT.dist(nu=nu, mu=xtm1, sigma=sigma, shape=size)
-                return x, collect_default_updates([x])
+            def step(xtm1, rng, nu, sigma):
+                next_rng, x = StudentT.dist(
+                    nu=nu, mu=xtm1, sigma=sigma, shape=size, rng=rng
+                ).owner.outputs
+                return x, next_rng
 
-            xs, _ = scan(
+            xs, _next_rng = scan(
                 fn=step,
-                outputs_info=pt.zeros(size),
+                outputs_info=[pt.zeros(size), rng],
                 non_sequences=[nu, sigma],
                 n_steps=steps,
+                return_updates=False,
             )
 
             # Logprob inference cannot be derived yet  https://github.com/pymc-devs/pymc/issues/6360
@@ -667,13 +679,17 @@ def step(s):
                 traffic = s + innov
                 return traffic, {innov.owner.inputs[0]: innov.owner.outputs[0]}
 
-            rv_seq, _ = pytensor.scan(
-                fn=step,
-                sequences=[seq],
-                outputs_info=[None],
-                n_steps=n_steps,
-                strict=True,
-            )
+            with pytest.warns(DeprecationWarning, match="Scan return signature will change"):
+                rv_seq, _ = pytensor.scan(
+                    fn=step,
+                    sequences=[seq],
+                    outputs_info=[None],
+                    n_steps=n_steps,
+                    strict=True,
+                    # There's a bug in the ordering of outputs when there's a mapped `None` output
+                    # We have to stick with the deprecated API for now
+                    return_updates=True,
+                )
             return rv_seq
 
         def normal_shifted(mu, size):