[mutable-arrays] remat discharge rule

Author: mattjj

jax/_src/core.py

@@ -3080,6 +3080,11 @@ def write(v: Var, a: AvalQDD) -> None:
                          f"from source: {src}"])
       raise JaxprTypeError(msg, eqn_idx) from None
 
+  # Check no returned refs  TODO(mattjj); improve this error message
+  from jax._src.state.types import AbstractRef  # pytype: disable=import-error
+  for v in jaxpr.outvars:
+    if isinstance(v.aval, AbstractRef): raise TypeError("returned ref")
+
   # TODO(mattjj): include output type annotation on jaxpr and check it here
   foreach(read, jaxpr.outvars)
 

jax/_src/state/discharge.py

@@ -22,6 +22,7 @@
 from typing import Any, Protocol, TypeVar
 
 from jax._src import ad_util
+from jax._src import ad_checkpoint
 from jax._src import api_util
 from jax._src import core
 from jax._src import linear_util as lu
@@ -1179,3 +1180,17 @@ def _pjit_state_discharge_rule(
   sentinel = object()
   assert next(ref_vals_iter, sentinel) is sentinel
   return new_invals, out_vals
+
+
+@register_discharge_rule(ad_checkpoint.remat_p)
+def _remat_state_discharge_rule(
+    in_avals, out_avals, *args, jaxpr, **params):
+  discharged_jaxpr, () = discharge_state(jaxpr, [])
+  out_vals_ref_vals = ad_checkpoint.remat_p.bind(
+      *args, jaxpr=discharged_jaxpr, **params)
+  out_vals, ref_vals = split_list(out_vals_ref_vals, [len(jaxpr.outvars)])
+  ref_vals_ = iter(ref_vals)
+  new_invals = [next(ref_vals_) if isinstance(a, AbstractRef) else None
+                for a in in_avals]
+  assert next(ref_vals_, None) is None
+  return new_invals, out_vals

tests/mutable_array_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from functools import partial
+import itertools as it
 import numpy as np
 import jax
 from jax._src import core
@@ -367,14 +368,16 @@ def loss(x, y):
 
   @parameterized.parameters([False, True])
   def test_custom_vjp_grad_stats_plumbing_basic(self, jit):
-    @jax.jit
     def primal(grads_ref, x):  # note: jit-abstracted!
       x = jnp.sin(x)
       x = stash_grads(grads_ref, x)
       x = jnp.sin(x)
       x = stash_grads(grads_ref, x)  # ignored, order-preserved
       return x
 
+    if jit:
+      primal = jax.jit(primal)
+
     @jax.custom_vjp
     def stash_grads(grads_ref, x):
       return x
@@ -389,18 +392,22 @@ def stash_grads_bwd(grads_ref, g):
     jax.grad(primal, 1)(grads_ref, jnp.float32(1.0))
     self.assertAllClose(grads_ref[...], jnp.cos(jnp.sin(1.)), check_dtypes=False)
 
-  @parameterized.parameters([False, True])
-  def test_custom_vjp_grad_stats_plumbing_scan(self, jit):
-    @jax.jit
+  @parameterized.parameters(it.product([False, True], repeat=2))
+  def test_custom_vjp_grad_stats_plumbing_scan(self, jit, remat):
     def primal(grads_ref, x):  # note: jit-abstracted!
       def body(x, _):
         x = jnp.sin(x)
         x = stash_grads(grads_ref, x)
         x = jnp.sin(x)
         return x, ()
+      if remat:
+        body = jax.remat(body)
       x, () = jax.lax.scan(body, x, None, length=1)
       return x
 
+    if jit:
+      primal = jax.jit(primal)
+
     @jax.custom_vjp
     def stash_grads(grads_ref, x):
       return x
@@ -415,6 +422,35 @@ def stash_grads_bwd(grads_ref, g):
     jax.grad(primal, argnums=1)(grads_ref, jnp.float32(1.0))
     self.assertAllClose(grads_ref[...], jnp.cos(jnp.sin(1.)), check_dtypes=False)
 
+  # TODO(mattjj,dougalm): this errors, which may or may not be desired behavior
+  # def test_remat_basic(self):
+  #   @jax.remat
+  #   def f(x_ref, y):
+  #     x_ref[...] += 1
+  #     return y
+
+  #   x_ref = core.mutable_array(0)
+  #   jax.grad(f, 1)(x_ref, 3.14)
+
+  def test_remat_grad_stats_plumbing_basic(self):
+    @jax.remat
+    def f(x_ref, y):
+      stash_grads(x_ref, y)
+      return y
+
+    @jax.custom_vjp
+    def stash_grads(grads_ref, x):
+      return x
+    def stash_grads_fwd(grads_ref, x):
+      return x, grads_ref
+    def stash_grads_bwd(grads_ref, g):
+      grads_ref[...] = g
+      return None, g
+    stash_grads.defvjp(stash_grads_fwd, stash_grads_bwd)
+
+    x_ref = core.mutable_array(0)
+    jax.grad(f, 1)(x_ref, 3.14)
+
 
 @jtu.with_config(jax_mutable_array_checks=True)
 class MutableArrayErrorsTest(jtu.JaxTestCase):