Merge pull request #33924 from mattjj:while-loop-internal-ref-effect

PiperOrigin-RevId: 843789461

Author: Google-ML-Automation

jax/_src/ad_checkpoint.py

@@ -845,7 +845,8 @@ def remat_dce(used_outputs: list[bool], eqn: core.JaxprEqn
 pe.dce_rules[remat_p] = remat_dce
 
 def _has_effects(effects) -> bool:
-  return bool({e for e in effects if not isinstance(e, core.NamedAxisEffect)})
+  not_really_effects = (core.NamedAxisEffect, core.InternalMutableArrayEffect)
+  return any(not isinstance(e, not_really_effects) for e in effects)
 
 
 def remat_expansion(

jax/_src/interpreters/partial_eval.py

@@ -1050,7 +1050,8 @@ def ensure_instantiated(inst: bool, x: Atom) -> Atom:
     return x
 
   def has_effects(effects) -> bool:
-    return bool({e for e in effects if not isinstance(e, core.NamedAxisEffect)})
+    not_really_effects = (core.NamedAxisEffect, core.InternalMutableArrayEffect)
+    return any(not isinstance(e, not_really_effects) for e in effects)
 
   known_eqns, staged_eqns = [], []
   foreach(write, in_unknowns, in_inst, jaxpr.invars)

tests/mutable_array_test.py

@@ -1053,6 +1053,34 @@ def f(x):
 
     f(3.)  # don't crash
 
+  def test_remat_while_loop_residuals(self):
+    @jax.custom_vjp
+    def ra2a(x):
+      return jax.freeze(jax.new_ref(x))
+
+    def ra2a_fwd(x):
+      o = ra2a(x)
+      return o, ()
+
+    def ra2a_bwd(res, g):
+      return (ra2a(g),)
+
+    ra2a.defvjp(ra2a_fwd, ra2a_bwd)
+
+    @jax.jit
+    @jax.remat
+    def f(x):
+
+      def g(x):
+        def body(carry):
+          i, x = carry
+          x = ra2a(x)
+          return i + 1, x
+        return jax.lax.while_loop(lambda x: x[0] < 5, body, (0, x))[1]
+      return g(x)
+
+    jax.linearize(f, 5.)  # don't crash
+
 
 @jtu.with_config(jax_mutable_array_checks=True)
 class MutableArrayErrorsTest(jtu.JaxTestCase):