[mutable-arrays] ignroe InternalMutableArrayEffect in partial_eval has_effects

It's not really an effect from the point of view of partial eval; like
NamedAxisEffect, it's just a bit of info we're sneaking into the effect system
as a handy way of tracking it. But it should be ignored for the purpose of
deciding "is this code effectful".

(We're hoping that InternalMutableArrayEffect, and the effect system more
generally, and partial eval, are not long for this world...)

This is a follow-up on #33906

Co-authored-by: Sharad Vikram  <sharadmv@google.com>

Author: mattjj

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

@@ -1051,7 +1051,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):