Fix JAX Scan for output ndim > 1

pytensor/link/jax/dispatch/scan.py

@@ -42,6 +42,13 @@ def scan(*outer_inputs):
             op.outer_non_seqs(outer_inputs),
         )  # JAX `init`
 
+        ndim_out_core = [getattr(x, "ndim", 0) for x in op.fgraph.outputs]
+        # ndim_in_core = [getattr(x, "ndim", 0) for x in mit_sot_init + sit_sot_init]
+
+        add_batchdim_flags = [dims_out > 0 for dims_out in ndim_out_core] + [
+            False
+        ] * op.info.n_nit_sot
+
         def jax_args_to_inner_func_args(carry, x):
             """Convert JAX scan arguments into format expected by scan_inner_func.
 
@@ -151,11 +158,18 @@ def get_partial_traces(traces):
                 + op.outer_nitsot(outer_inputs)
             )
             partial_traces = []
-            for init_state, trace, buffer in zip(init_states, traces, buffers):
+            for init_state, trace, add_batchdim, buffer in zip(
+                init_states, traces, add_batchdim_flags, buffers
+            ):
                 if init_state is not None:
                     # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
+                    if add_batchdim:
+                        init_state = jnp.expand_dims(init_state, 0)
                     full_trace = jnp.concatenate(
-                        [jnp.atleast_1d(init_state), jnp.atleast_1d(trace)],
+                        [
+                            jnp.atleast_1d(init_state),
+                            jnp.atleast_1d(trace),
+                        ],
                         axis=0,
                     )
                     buffer_size = buffer.shape[0]

tests/link/jax/test_scan.py

@@ -13,7 +13,7 @@
 from pytensor.scan.op import Scan
 from pytensor.tensor import random
 from pytensor.tensor.math import gammaln, log
-from pytensor.tensor.type import lscalar, scalar, vector
+from pytensor.tensor.type import dmatrix, dvector, lscalar, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
@@ -317,3 +317,100 @@ def input_step_fn(y_tm1, y_tm3, a):
 
     test_input_vals = [np.array(10.0).astype(config.floatX)]
     compare_jax_and_py(out_fg, test_input_vals)
+
+
+@pytest.mark.parametrize("x0_func", [dvector, dmatrix])
+@pytest.mark.parametrize("A_func", [dmatrix, dmatrix])
+def test_nd_scan_sit_sot(x0_func, A_func):
+    x0 = x0_func("x0")
+    A = A_func("A")
+
+    n_steps = 3
+    k = 3
+
+    # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
+    xs, _ = scan(
+        lambda X, A: A @ X,
+        non_sequences=[A],
+        outputs_info=[x0],
+        n_steps=n_steps,
+        mode=get_mode("JAX"),
+    )
+
+    x0_val = np.arange(k) if x0.ndim == 1 else np.diag(np.arange(k))
+    A_val = np.eye(k)
+
+    fg = FunctionGraph([x0, A], [xs])
+    test_input_vals = [x0_val, A_val]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+def test_nd_scan_sit_sot_with_seq():
+    x = dmatrix("x0")
+    A = dmatrix("A")
+
+    n_steps = 3
+    k = 3
+
+    # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
+    xs, _ = scan(
+        lambda X, A: A @ X,
+        non_sequences=[A],
+        sequences=[x],
+        n_steps=n_steps,
+        mode=get_mode("JAX"),
+    )
+
+    x_val = np.tile(np.arange(k), n_steps).reshape(n_steps, k)
+    A_val = np.eye(k)
+
+    fg = FunctionGraph([x, A], [xs])
+    test_input_vals = [x_val, A_val]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+def test_nd_scan_mit_sot():
+    x0 = dmatrix("x0")
+    A = dmatrix("A")
+    B = dmatrix("B")
+
+    # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
+    xs, _ = scan(
+        lambda xtm3, xtm1, A, B: A @ xtm3 + B @ xtm1,
+        outputs_info=[{"initial": x0, "taps": [-3, -1]}],
+        non_sequences=[A, B],
+        n_steps=10,
+        mode=get_mode("JAX"),
+    )
+
+    fg = FunctionGraph([x0, A, B], [xs])
+    x0_val = np.r_[[np.arange(3).tolist()] * 3]
+    A_val = np.eye(3)
+    B_val = np.eye(3)
+
+    test_input_vals = [x0_val, A_val, B_val]
+    compare_jax_and_py(fg, test_input_vals)
+
+
+def test_nd_scan_sit_sot_with_carry():
+    x0 = dvector("x0")
+    A = dmatrix("A")
+
+    def step(x, A):
+        return A @ x, x.sum()
+
+    # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
+    xs, _ = scan(
+        step,
+        outputs_info=[x0, None],
+        non_sequences=[A],
+        n_steps=10,
+        mode=get_mode("JAX"),
+    )
+
+    fg = FunctionGraph([x0, A], xs)
+    x0_val = np.arange(3)
+    A_val = np.eye(3)
+
+    test_input_vals = [x0_val, A_val]
+    compare_jax_and_py(fg, test_input_vals)