Switch to new Scan API

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

@@ -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")

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

@@ -356,19 +356,21 @@ def test_custom_dist_default_support_point(self, dist_params, size, expected, di
         assert_support_point_is_expected(model, expected)
 
     def test_custom_dist_default_support_point_scan(self):
-        def scan_step(left, right):
-            x = Uniform.dist(left, right)
-            x_update = collect_default_updates([x])
-            return x, x_update
+        def scan_step(left, right, rng):
+            next_rng, x = Uniform.dist(left, right, rng=rng)
+            return x, next_rng
 
         def dist(size):
-            xs, updates = scan(
+            rng = pytensor.shared(np.random.default_rng())
+            xs, _next_rng = scan(
                 fn=scan_step,
+                outputs_info=[rng],
                 sequences=[
                     pt.as_tensor_variable(np.array([-4, -3])),
                     pt.as_tensor_variable(np.array([-2, -1])),
                 ],
                 name="xs",
+                return_updates=False,
             )
             return xs
 
@@ -377,17 +379,21 @@ def dist(size):
         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
 
@@ -662,17 +668,22 @@ def test_chained_custom_dist_bug(self):
         batch = 2
 
         def scan_dist(seq, n_steps, size):
-            def step(s):
-                innov = Normal.dist()
+            rng = pytensor.shared(np.random.default_rng(123))
+
+            def step(s, rng):
+                next_rng, innov = Normal.dist(rng=rng).owner.outputs
                 traffic = s + innov
-                return traffic, {innov.owner.inputs[0]: innov.owner.outputs[0]}
+                return traffic, next_rng
 
-            rv_seq, _ = pytensor.scan(
+            rv_seq, _next_rng = pytensor.scan(
                 fn=step,
                 sequences=[seq],
-                outputs_info=[None],
+                # There's a bug in the pytensor helper that's flipping the two returned outputs
+                # outputs_info=[None, rng],
+                outputs_info=[rng, None],
                 n_steps=n_steps,
                 strict=True,
+                return_updates=False,
             )
             return rv_seq