Fix Scan JAX dispatcher

Author: ricardoV94

pytensor/link/jax/dispatch/scan.py

@@ -1,20 +1,36 @@
 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)
+    info = op.info
+
+    if info.as_while:
+        raise NotImplementedError("While Scan cannot yet be converted to JAX")
+
+    if info.n_mit_mot:
+        raise NotImplementedError(
+            "Scan with MIT-MOT (gradients of scan) cannot yet be converted to JAX"
+        )
+
+    # Optimize inner graph
+    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,31 +45,23 @@ def scan(*outer_inputs):
         n_steps = scan_args.n_steps
         seqs = scan_args.outer_in_seqs
 
-        # TODO: mit_mots
-        mit_mot_in_slices = []
-
         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]
 
         init_carry = (
-            mit_mot_in_slices,
+            [],  # mit_mot_in_slices
             mit_sot_in_slices,
             sit_sot_in_slices,
             scan_args.outer_in_shared,
             scan_args.outer_in_non_seqs,
         )
 
         def jax_args_to_inner_scan(op, carry, x):
-            # `carry` contains all inner-output taps, non_seqs, and shared
-            # terms
+            # `carry` contains all inner-output taps, non_seqs, and shared terms
             (
                 inner_in_mit_mot,
                 inner_in_mit_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,
@@ -103,57 +112,131 @@ def inner_scan_outs_to_jax_outs(
                 inner_in_non_seqs,
             ) = old_carry
 
-            def update_mit_sot(mit_sot, new_val):
-                return jnp.concatenate([mit_sot[1:], new_val[None, ...]], axis=0)
+            inner_out_mit_sot = inner_scan_outs[
+                info.n_mit_mot : info.n_mit_mot + info.n_mit_sot
+            ]
+            inner_in_mit_sot_new = []
+            if inner_in_mit_sot:
+                # Replace the oldest tap by the newest value
+                inner_in_mit_sot_new = [
+                    jnp.concatenate([old_mit_sot[1:], new_val[None, ...]], axis=0)
+                    for old_mit_sot, new_val in zip(
+                        inner_in_mit_sot,
+                        inner_out_mit_sot,
+                    )
+                ]
+
+            inner_out_sit_sot = inner_in_sit_sot_new = inner_scan_outs[
+                info.n_mit_mot
+                + info.n_mit_sot : info.n_mit_mot
+                + info.n_mit_sot
+                + info.n_sit_sot
+            ]
 
-            inner_out_mit_sot = [
-                update_mit_sot(mit_sot, new_val)
-                for mit_sot, new_val in zip(inner_in_mit_sot, inner_scan_outs)
+            inner_out_nit_sot = inner_scan_outs[
+                info.n_mit_mot
+                + info.n_mit_sot
+                + info.n_sit_sot : info.n_mit_mot
+                + info.n_mit_sot
+                + info.n_sit_sot
+                + info.n_nit_sot :
             ]
 
-            # This should contain all inner-output taps, non_seqs, and shared
-            # terms
-            if not inner_in_sit_sot:
-                inner_out_sit_sot = []
-            else:
-                inner_out_sit_sot = inner_scan_outs
+            inner_in_shared_new = inner_in_shared
+            if info.n_shared_outs:
+                # Replace old shared inputs by new shared outputs
+                new_inner_out_shared = inner_scan_outs[
+                    info.n_mit_mot + info.n_mit_sot + info.n_sit_sot + info.n_nit_sot :
+                ]
+                inner_in_shared_new[: info.n_shared_outs] = new_inner_out_shared
+
             new_carry = (
-                inner_in_mit_mot,
-                inner_out_mit_sot,
-                inner_out_sit_sot,
-                inner_in_shared,
+                [],  # MIT-MOT
+                inner_in_mit_sot_new,
+                inner_in_sit_sot_new,
+                inner_in_shared_new,
                 inner_in_non_seqs,
             )
 
-            return new_carry
+            # Shared variables and non_seqs are not traced
+            new_scan = sum(
+                [
+                    [],  # MIT-MOT
+                    inner_out_mit_sot,
+                    inner_out_sit_sot,
+                    inner_out_nit_sot,
+                ],
+                [],
+            )
+
+            return new_carry, new_scan
 
         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))
-            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)
-
-        # 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)
-            ]
+            inner_scan_outs = list(scan_inner_func(*inner_args))
+            new_carry, new_scan_outs = inner_scan_outs_to_jax_outs(
+                op, carry, inner_scan_outs
+            )
+            return new_carry, new_scan_outs
+
+        last_state, scan_traces = jax.lax.scan(
+            jax_inner_func, init_carry, seqs, length=n_steps
+        )
+
+        def get_partial_traces(scan_traces):
+            # 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.
+            init_states = (
+                mit_sot_in_slices
+                + sit_sot_in_slices
+                + [None] * len(scan_args.outer_in_nit_sot)
+            )
+            buffers = (
+                scan_args.outer_in_mit_sot
+                + scan_args.outer_in_sit_sot
+                + scan_args.outer_in_nit_sot
+            )
+
+            partial_scan_traces = []
+            for init_state, scan_trace, buffer in zip(
+                init_states, scan_traces, buffers
+            ):
+                if init_state is not None:
+                    # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
+                    full_scan_trace = jnp.concatenate(
+                        [jnp.atleast_1d(init_state), jnp.atleast_1d(scan_trace)],
+                        axis=0,
+                    )
+                    partial_scan_trace = full_scan_trace[-buffer.shape[0] :]
+                else:
+                    # NIT-SOT: Buffer is just the number of entries that should be returned
+                    partial_scan_trace = jnp.atleast_1d(scan_trace)[-buffer:]
+                partial_scan_traces.append(partial_scan_trace)
+
+            return partial_scan_traces
+
+        def get_shared_outs(last_state):
+            # Select the last state of shared_outs, these outputs are not traced
+            if not info.n_shared_outs:
+                return []
+
+            (
+                inner_out_mit_mot,
+                inner_out_mit_sot,
+                inner_out_sit_sot,
+                inner_out_shared,
+                inner_in_non_seqs,
+            ) = last_state
+
+            # TODO: Check if a shared variable that is not an output shows up here or in non-seqs
+            shared_outs = inner_out_shared[: info.n_shared_outs]
+            return list(shared_outs)
+
+        scan_outs_final = get_partial_traces(scan_traces) + get_shared_outs(last_state)
 
-        if len(scan_out_final) == 1:
-            scan_out_final = scan_out_final[0]
-        return scan_out_final
+        if len(scan_outs_final) == 1:
+            scan_outs_final = scan_outs_final[0]
+        return scan_outs_final
 
     return scan