Merge pull request jax-ml#25369 from jax-ml:mutable-arrays-ad

PiperOrigin-RevId: 704685653

Author: Google-ML-Automation

jax/_src/core.py

@@ -431,6 +431,8 @@ class Primitive:
   call_primitive: bool = False
   # set for map primitives processed in final style.
   map_primitive: bool = False
+  # set for ref primitives
+  ref_primitive: bool = False
 
   def __init__(self, name: str):
     self.name = name
@@ -1882,6 +1884,7 @@ def __repr__(self) -> str: return 'Mutable' + repr(self[...])
 def mutable_array(init_val):
   return mutable_array_p.bind(init_val)
 mutable_array_p = Primitive('mutable_array')
+mutable_array_p.ref_primitive = True
 
 class InternalMutableArrayEffect(effects.Effect):
   pass
@@ -1899,6 +1902,18 @@ def _mutable_array_impl(init_val):
   aval = get_aval(init_val)
   return MutableArray(AbstractRef(aval), init_val)
 
+def freeze(ref):
+  return freeze_p.bind(ref)
+freeze_p = Primitive('freeze')
+freeze_p.ref_primitive = True
+
+@freeze_p.def_effectful_abstract_eval
+def freeze_abstract_eval(ref_aval):
+  return ref_aval.inner_aval, {internal_mutable_array_effect}
+
+@freeze_p.def_impl
+def _freeze_impl(ref):
+  return ref[()]
 
 class AbstractToken(AbstractValue):
   def str_short(self, short_dtypes=False): return 'Tok'
@@ -2516,10 +2531,11 @@ def write(v: Var, a: AbstractValue) -> None:
 
       # Check the computed effect type matches the eqn's annotation, and is
       # included in the jaxpr's annotation.
-      if prim is mutable_array_p:
-        outvar, = eqn.outvars
-        in_idx[outvar] = None  # type: ignore
-        mut_arrays.add(outvar)
+      if prim.ref_primitive:
+        if prim is mutable_array_p:
+          outvar, = eqn.outvars
+          in_idx[outvar] = None  # type: ignore
+          mut_arrays.add(outvar)
       if eqn.effects != eqn_effects:
         raise JaxprTypeError("Inferred effects do not match equation effects. "
                              f"Equation effects: {eqn.effects}. "

jax/_src/interpreters/ad.py

@@ -267,6 +267,20 @@ def write_primal(v, val):
   with ctx:
     map(partial(write_cotangent, 'outvars'), jaxpr.outvars, cotangents_in)
     for eqn in jaxpr.eqns[::-1]:
+      if eqn.primitive.ref_primitive:
+        if eqn.primitive is core.mutable_array_p:
+          val_var, = eqn.invars
+          ref_var, = eqn.outvars
+          ref = read_primal(ref_var)
+          ct_out = core.freeze(ref)
+          write_cotangent(eqn.primitive, val_var, ct_out)
+        elif eqn.primitive is core.freeze_p:
+          val_var, = eqn.outvars
+          ref_var, = eqn.invars
+          ct_in = instantiate_zeros(read_cotangent(val_var))
+          write_primal(ref_var, core.mutable_array(ct_in))
+        continue
+
       invals = map(read_primal, eqn.invars)
       if eqn.primitive.multiple_results:
         cts_in = map(read_cotangent, eqn.outvars)

jax/_src/interpreters/partial_eval.py

@@ -1009,21 +1009,25 @@ def partial_eval_jaxpr_stateful(
     in_inst: bool | Sequence[bool],
     ensure_out_unknowns: bool | Sequence[bool],
     ensure_out_inst: bool | Sequence[bool],
-    saveable: Callable[..., RematCases_],
+    saveable: Callable[..., RematCases_] | None,
   ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int, int]:
   if type(in_inst) is bool:
     in_inst = (in_inst,) * len(jaxpr.invars)
   if type(ensure_out_unknowns) is bool:
     ensure_out_unknowns = (ensure_out_unknowns,) * len(jaxpr.outvars)
   if type(ensure_out_inst) is bool:
     ensure_out_inst = (ensure_out_inst,) * len(jaxpr.outvars)
+  if saveable is None:
+    saveable = everything_saveable
   jaxpr_known, jaxpr_staged, out_unknowns, out_inst, num_res, num_res_ref = \
       _partial_eval_jaxpr_custom_cached(jaxpr, tuple(in_unknowns),
                                         tuple(in_inst),
                                         tuple(ensure_out_unknowns),
                                         tuple(ensure_out_inst), saveable)
   return jaxpr_known, jaxpr_staged, out_unknowns, out_inst, num_res, num_res_ref
 
+everything_saveable = lambda *_, **__: True
+
 @weakref_lru_cache
 def _partial_eval_jaxpr_custom_cached(
     jaxpr: Jaxpr,

jax/_src/lax/lax.py

@@ -1730,6 +1730,9 @@ def zeros_like_abstract_ref(aval: state.AbstractRef) -> core.MutableArray:
   val = ad_util.zeros_like_aval(aval.inner_aval)
   return core.mutable_array(val)
 
+# TODO(dougalm): this is nonsense but it's here because in places like
+# custom_vjp we assume that all arguments have tangent spaces. We could have
+# a distinct NotATangentType value instead.
 ad_util.aval_zeros_likers[state.AbstractRef] = zeros_like_abstract_ref  # type: ignore
 
 def iota(dtype: DTypeLike, size: int) -> Array:

jax/_src/pjit.py

@@ -2156,8 +2156,18 @@ def _pjit_partial_eval(trace, *in_tracers,
 
   known_ins = tuple(pv.is_known() for pv in in_pvals)
   unknown_ins = tuple(not k for k in known_ins)
-  known_jaxpr, unknown_jaxpr, unknown_outs, res_avals = \
-      pe.partial_eval_jaxpr_nounits(jaxpr, unknown_ins, instantiate=False)
+  if any(isinstance(e, (RefEffect, core.InternalMutableArrayEffect))
+         for e in jaxpr.effects):
+    known_jaxpr_, unknown_jaxpr_, unknown_outs, _, num_res_val, num_res_ref = \
+        pe.partial_eval_jaxpr_stateful(jaxpr.jaxpr, unknown_ins, unknown_ins,
+                                       False, False, None)
+    if num_res_ref: raise NotImplementedError
+    known_jaxpr = pe.ClosedJaxpr(known_jaxpr_, jaxpr.consts)
+    unknown_jaxpr = pe.ClosedJaxpr(unknown_jaxpr_, jaxpr.consts)
+    res_avals = unknown_jaxpr.in_avals[:num_res_val]
+  else:
+    known_jaxpr, unknown_jaxpr, unknown_outs, res_avals = \
+        pe.partial_eval_jaxpr_nounits(jaxpr, unknown_ins, instantiate=False)
   unknown_outs = tuple(unknown_outs)
   known_outs = tuple(not uk for uk in unknown_outs)
   num_residuals = len(res_avals)

jax/_src/state/discharge.py

@@ -153,6 +153,10 @@ def _eval_jaxpr_discharge_state(
       [invar], [outvar] = eqn.invars, eqn.outvars
       ans = env.read(invar)
       refs_to_discharge.add(id(outvar.aval))
+    elif eqn.primitive is core.freeze_p:
+      [invar], [outvar] = eqn.invars, eqn.outvars
+      ans = env.read(invar)
+      refs_to_discharge.remove(id(invar.aval))
     elif (any(should_discharge)
           or core.internal_mutable_array_effect in eqn.effects
       ):

jax/_src/state/primitives.py

@@ -654,3 +654,8 @@ def _broadcast_to_abstract_eval(aval, *, shape):
 mlir.register_lowering(
     broadcast_to_p, mlir.lower_fun(_broadcast_to_impl, False)
 )
+
+# === AD rules for mutable arrays ===
+
+ad.defjvp(core.mutable_array_p, lambda g, _: core.mutable_array(g))
+ad.defjvp(core.freeze_p, lambda g, _: core.freeze(g))

tests/mutable_array_test.py

@@ -232,5 +232,18 @@ def f():
     x = f()
     self.assertArraysEqual(x, jnp.zeros(8))
 
+  def test_grad_mutable_array(self):
+    @jax.jit
+    def f(x):
+      x_ = core.mutable_array(x)
+      x_[()] = x_[()] + x_[()]
+      y = core.freeze(x_)
+      return y
+
+    ans = jax.grad(f)(1.)
+    expected = 2.0
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())