Fix Scan JAX dispatcher

Author: ricardoV94

pytensor/link/jax/dispatch/scan.py

@@ -1,20 +1,31 @@
 import jax
 import jax.numpy as jnp
 
-from pytensor.graph.fg import FunctionGraph
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.scan.op import Scan
 from pytensor.scan.utils import ScanArgs
 
 
 @jax_funcify.register(Scan)
 def jax_funcify_Scan(op, **kwargs):
-    inner_fg = FunctionGraph(op.inputs, op.outputs)
-    jax_at_inner_func = jax_funcify(inner_fg, **kwargs)
+    # TODO: Raise NotImplementedError if While scan
+
+    # Apply inner rewrites
+    # TODO: Not sure this is the right place to do this, should we have a rewrite that
+    #  explicitly triggers the optimization of the inner graphs of Scan?
+    #  The C-code defers it to the make_thunk phase
+    fgraph = op.fgraph.clone()
+    rewriter = op.mode_instance.optimizer
+    rewriter(fgraph)
+    scan_inner_func = jax_funcify(fgraph, **kwargs)
 
     def scan(*outer_inputs):
         scan_args = ScanArgs(
-            list(outer_inputs), [None] * op.info.n_outs, op.inputs, op.outputs, op.info
+            list(outer_inputs),
+            [None] * len(op.inner_outputs),
+            op.inner_inputs,
+            op.inner_outputs,
+            op.info,
         )
 
         # `outer_inputs` is a list with the following composite form:
@@ -29,16 +40,13 @@ def scan(*outer_inputs):
         n_steps = scan_args.n_steps
         seqs = scan_args.outer_in_seqs
 
-        # TODO: mit_mots
         mit_mot_in_slices = []
+        if scan_args.outer_in_mit_mot:
+            raise NotImplementedError("JAX Scan with MIT-MOT not supported yet.")
 
         mit_sot_in_slices = []
         for tap, seq in zip(scan_args.mit_sot_in_slices, scan_args.outer_in_mit_sot):
-            neg_taps = [abs(t) for t in tap if t < 0]
-            pos_taps = [abs(t) for t in tap if t > 0]
-            max_neg = max(neg_taps) if neg_taps else 0
-            max_pos = max(pos_taps) if pos_taps else 0
-            init_slice = seq[: max_neg + max_pos]
+            init_slice = seq[: abs(min(tap))]
             mit_sot_in_slices.append(init_slice)
 
         sit_sot_in_slices = [seq[0] for seq in scan_args.outer_in_sit_sot]
@@ -76,6 +84,7 @@ def jax_args_to_inner_scan(op, carry, x):
             for array, index in zip(inner_in_mit_sot, scan_args.mit_sot_in_slices):
                 inner_in_mit_sot_flatten.extend(array[jnp.array(index)])
 
+            # Concatenate lists
             inner_scan_inputs = sum(
                 [
                     inner_in_seqs,
@@ -116,9 +125,13 @@ def update_mit_sot(mit_sot, new_val):
             if not inner_in_sit_sot:
                 inner_out_sit_sot = []
             else:
-                inner_out_sit_sot = inner_scan_outs
+                inner_out_sit_sot = inner_scan_outs[
+                    len(inner_in_mit_sot) : len(inner_in_mit_sot)
+                    + len(inner_in_sit_sot)
+                ]
+
             new_carry = (
-                inner_in_mit_mot,
+                inner_in_mit_mot,  # Just an empty list, we raise earlier if there are any MIT-MOT
                 inner_out_mit_sot,
                 inner_out_sit_sot,
                 inner_in_shared,
@@ -129,28 +142,39 @@ def update_mit_sot(mit_sot, new_val):
 
         def jax_inner_func(carry, x):
             inner_args = jax_args_to_inner_scan(op, carry, x)
-            inner_scan_outs = list(jax_at_inner_func(*inner_args))
+            inner_scan_outs = list(scan_inner_func(*inner_args))
             new_carry = inner_scan_outs_to_jax_outs(op, carry, inner_scan_outs)
             return new_carry, inner_scan_outs
 
-        _, scan_out = jax.lax.scan(jax_inner_func, init_carry, seqs, length=n_steps)
+        _, scan_outs = jax.lax.scan(jax_inner_func, init_carry, seqs, length=n_steps)
 
         # We need to prepend the initial values so that the JAX output will
         # match the raw `Scan` `Op` output and, thus, work with a downstream
         # `Subtensor` `Op` introduced by the `scan` helper function.
-        def append_scan_out(scan_in_part, scan_out_part):
-            return jnp.concatenate([scan_in_part[:-n_steps], scan_out_part], axis=0)
-
-        if scan_args.outer_in_mit_sot:
-            scan_out_final = [
-                append_scan_out(init, out)
-                for init, out in zip(scan_args.outer_in_mit_sot, scan_out)
-            ]
-        elif scan_args.outer_in_sit_sot:
-            scan_out_final = [
-                append_scan_out(init, out)
-                for init, out in zip(scan_args.outer_in_sit_sot, scan_out)
-            ]
+        scan_out_final = []
+        for init, scan_out, buffer in zip(
+            mit_sot_in_slices
+            + sit_sot_in_slices
+            + [None] * len(scan_args.outer_in_nit_sot),
+            scan_outs,
+            scan_args.outer_in_mit_sot
+            + scan_args.outer_in_sit_sot
+            + scan_args.outer_in_nit_sot,
+        ):
+            if init is not None:
+                # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
+                full_scan_out = jnp.concatenate(
+                    [
+                        jnp.atleast_1d(init),
+                        jnp.atleast_1d(scan_out),
+                    ],
+                    axis=0,
+                )
+                partial_scan_out = full_scan_out[-buffer.shape[0] :]
+            else:
+                # NIT-SOT: Buffer is just the number of entries that should be returned
+                partial_scan_out = jnp.atleast_1d(scan_out)[-buffer:]
+            scan_out_final.append(partial_scan_out)
 
         if len(scan_out_final) == 1:
             scan_out_final = scan_out_final[0]

tests/link/jax/test_scan.py

@@ -1,24 +1,70 @@
 import numpy as np
 import pytest
-from packaging.version import parse as version_parse
 
 import pytensor.tensor as at
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
 from pytensor.scan.basic import scan
 from pytensor.tensor.math import gammaln, log
-from pytensor.tensor.type import ivector, lscalar, scalar
+from pytensor.tensor.type import lscalar, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
 jax = pytest.importorskip("jax")
 
 
-@pytest.mark.xfail(
-    version_parse(jax.__version__) >= version_parse("0.2.12"),
-    reason="Omnistaging cannot be disabled",
-)
-def test_jax_scan_multiple_output():
+@pytest.mark.parametrize("view", [None, (-1,), slice(-2, None, None)])
+def test_simple_sit_sot_scan(view):
+    x0 = at.scalar("x0", dtype="float64")
+    xs, _ = scan(
+        lambda xtm1: xtm1 + 1,
+        outputs_info=[x0],
+        n_steps=10,
+    )
+    if view:
+        xs = xs[view]
+    fg = FunctionGraph([x0], [xs])
+    test_input_vals = [np.e]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+@pytest.mark.parametrize("view", [None, (-1,), slice(-4, -1, None)])
+def test_simple_mit_sot_scan(view):
+    x0 = at.vector("x0", dtype="float64", shape=(3,))
+    xs, _ = scan(
+        lambda xtm3, xtm1: xtm3 + xtm1 + 1,
+        outputs_info=[{"initial": x0, "taps": [-3, -1]}],
+        n_steps=10,
+    )
+    if view:
+        xs = xs[view]
+    fg = FunctionGraph([x0], [xs])
+    test_input_vals = [np.full((3,), np.e)]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+@pytest.mark.parametrize("view", [None, (-2,), slice(None, None, 2)])
+def test_simple_nit_sot_scan(view):
+    rng = np.random.default_rng(seed=49)
+
+    xs = at.vector("x0", dtype="float64", shape=(10,))
+    s0 = at.zeros(())
+
+    # We need to have a recurring state, otherwise this simple scan would
+    # be rewritten as a simple Elemwise on xs
+    [_, ys], _ = scan(
+        lambda x, s: (s + 1, at.exp(x + s)),
+        outputs_info=[s0, None],
+        sequences=[xs],
+    )
+    if view:
+        ys = ys[view]
+    fg = FunctionGraph([xs], [ys])
+    test_input_vals = [rng.normal(size=10)]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+def test_SEIR_scan():
     """Test a scan implementation of a SEIR model.
 
     SEIR model definition:
@@ -38,8 +84,8 @@ def binom_log_prob(n, p, value):
         return binomln(n, value) + value * log(p) + (n - value) * log(1 - p)
 
     # sequences
-    at_C = ivector("C_t")
-    at_D = ivector("D_t")
+    at_C = vector("C_t", dtype="int32", shape=(8,))
+    at_D = vector("D_t", dtype="int32", shape=(8,))
     # outputs_info (initial conditions)
     st0 = lscalar("s_t0")
     et0 = lscalar("e_t0")
@@ -108,11 +154,7 @@ def seir_one_step(ct0, dt0, st0, et0, it0, logp_c, logp_d, beta, gamma, delta):
     compare_jax_and_py(out_fg, test_input_vals)
 
 
-@pytest.mark.xfail(
-    version_parse(jax.__version__) >= version_parse("0.2.12"),
-    reason="Omnistaging cannot be disabled",
-)
-def test_jax_scan_tap_output():
+def test_mitsot_with_nonseq_scan():
     a_at = scalar("a")
 
     def input_step_fn(y_tm1, y_tm3, a):