re-write ndim check, specify static sizes for tests

Author: jessegrabowski

pytensor/link/jax/dispatch/scan.py

@@ -42,13 +42,6 @@ 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.
 
@@ -158,20 +151,14 @@ def get_partial_traces(traces):
                 + op.outer_nitsot(outer_inputs)
             )
             partial_traces = []
-            for init_state, trace, add_batchdim, buffer in zip(
-                init_states, traces, add_batchdim_flags, buffers
-            ):
+            for init_state, trace, buffer in zip(init_states, traces, 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),
-                        ],
-                        axis=0,
+                    trace = jnp.atleast_1d(trace)
+                    init_state = jnp.expand_dims(
+                        init_state, range(trace.ndim - init_state.ndim)
                     )
+                    full_trace = jnp.concatenate([init_state, trace], axis=0)
                     buffer_size = buffer.shape[0]
                 else:
                     # NIT-SOT: Buffer is just the number of entries that should be returned

tests/link/jax/test_scan.py

@@ -346,12 +346,12 @@ def test_nd_scan_sit_sot(x0_func, A_func):
 
 
 def test_nd_scan_sit_sot_with_seq():
-    x = dmatrix("x0")
-    A = dmatrix("A")
-
     n_steps = 3
     k = 3
 
+    x = at.matrix("x0", shape=(n_steps, k))
+    A = at.matrix("A", shape=(k, k))
+
     # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
     xs, _ = scan(
         lambda X, A: A @ X,
@@ -370,9 +370,9 @@ def test_nd_scan_sit_sot_with_seq():
 
 
 def test_nd_scan_mit_sot():
-    x0 = dmatrix("x0")
-    A = dmatrix("A")
-    B = dmatrix("B")
+    x0 = at.matrix("x0", shape=(3, 3))
+    A = at.matrix("A", shape=(3, 3))
+    B = at.matrix("B", shape=(3, 3))
 
     # Must specify mode = JAX for the inner func to avoid a GEMM Op in the JAX graph
     xs, _ = scan(
@@ -385,6 +385,7 @@ def test_nd_scan_mit_sot():
 
     fg = FunctionGraph([x0, A, B], [xs])
     x0_val = np.r_[[np.arange(3).tolist()] * 3]
+    print(x0_val)
     A_val = np.eye(3)
     B_val = np.eye(3)
 
@@ -393,8 +394,8 @@ def test_nd_scan_mit_sot():
 
 
 def test_nd_scan_sit_sot_with_carry():
-    x0 = dvector("x0")
-    A = dmatrix("A")
+    x0 = at.vector("x0", shape=(3,))
+    A = at.matrix("A", shape=(3, 3))
 
     def step(x, A):
         return A @ x, x.sum()