Exclude backend incompatible rewrites in Scan dispatch

Author: ricardoV94

pytensor/link/jax/dispatch/scan.py

@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
 
+from pytensor.compile.mode import JAX
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.scan.op import Scan
 
@@ -17,8 +18,8 @@ def jax_funcify_Scan(op: Scan, **kwargs):
             "Scan with MIT-MOT (gradients of scan) cannot yet be converted to JAX"
         )
 
-    # Optimize inner graph
-    rewriter = op.mode_instance.optimizer
+    # Optimize inner graph (exclude any defalut rewrites that are incompatible with JAX mode)
+    rewriter = op.mode_instance.excluding(*JAX._optimizer.exclude).optimizer
     rewriter(op.fgraph)
     scan_inner_func = jax_funcify(op.fgraph, **kwargs)
 

pytensor/link/numba/dispatch/scan.py

@@ -5,6 +5,7 @@
 from numba import types
 from numba.extending import overload
 
+from pytensor.compile.mode import NUMBA
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
     create_arg_string,
@@ -58,7 +59,7 @@ def numba_funcify_Scan(op, node, **kwargs):
     # 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
-    rewriter = op.mode_instance.optimizer
+    rewriter = op.mode_instance.excluding(*NUMBA._optimizer.exclude).optimizer
     rewriter(op.fgraph)
 
     scan_inner_func = numba_basic.numba_njit(numba_funcify(op.fgraph))

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 dmatrix, dvector, lscalar, scalar, vector
+from pytensor.tensor.type import dmatrix, dvector, lscalar, matrix, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
 
 
@@ -418,3 +418,12 @@ def step(x, A):
 
     test_input_vals = [x0_val, A_val]
     compare_jax_and_py(fg, test_input_vals)
+
+
+def test_default_mode_excludes_incompatible_rewrites():
+    # See issue #426
+    A = matrix("A")
+    B = matrix("B")
+    out, _ = scan(lambda a, b: a @ b, outputs_info=[A], non_sequences=[B], n_steps=2)
+    fg = FunctionGraph([A, B], [out])
+    compare_jax_and_py(fg, [np.eye(3), np.eye(3)])