Merge pull request #33217 from gspschmid:gschmid/remat-ob-fix-consts

PiperOrigin-RevId: 831470125

Author: Google-ML-Automation

jax/_src/ad_checkpoint.py

@@ -363,7 +363,8 @@ def fun_remat(*args, **kwargs):
     in_avals = [core.shaped_abstractify(x) for x in args_flat]
     jaxpr, consts, out_tree = _trace_to_jaxpr(fun_, in_tree, tuple(in_avals), debug)
     if isinstance(prevent_cse, tuple):
-      cse = (*broadcast_prefix(prevent_cse, (args, kwargs) if kwargs else args),)
+      cse_args = (tuple(args), kwargs) if kwargs else tuple(args)
+      cse = (False,) * len(consts) + tuple(broadcast_prefix(prevent_cse, cse_args))
     else:
       cse = prevent_cse
     out_flat = remat_p.bind(

tests/api_test.py

@@ -7117,7 +7117,9 @@ def make_weight(i):
   def test_remat_partial_cse_prevention(self):
     @partial(jax.remat, prevent_cse=(False, True))
     def layer(W, x):
-      return x @ W
+      res = x @ W
+      res += jnp.array([1.0, 2.0, 3.0])  # ensure the jaxpr also contains a const
+      return res
 
     def net(Ws, x):
       for W in Ws: