Repository: google/jax

Issue:
Using `jax.jit` inside a function decorated by `jax.checkpoint` causes recompilation every time
Using a jitted function inside a function decorated by `jax.checkpoint` causes a lot of extra compilations even if the arguments still have the same shape. Calculating the gradient for such a function causes a memory leak in the long rung since all the compiled jitted functions seem to be stored in the memory. This can be observed by the high memory footprint of `backend_compile` which cannot be seen if the checkpointing is disabled.

A self-consistent example would be:
```python
import jax
import jax.numpy as jnp

@jax.jit
def f(a):
    return jnp.sum(a)

@jax.checkpoint
def g(a):
    return f(a)

arr = jnp.array([[1,2,3,4,5],[6,7,8,9,10]], dtype=float)

g_v_and_grad = jax.value_and_grad(g)

for i in range(3):
    working_arr = arr + i
    print(g_v_and_grad(working_arr))
```

Running the script with `env JAX_LOG_COMPILES=1` enabled one can observe:
```
WARNING:absl:Finished tracing + transforming prim_fun for jit in 0.0002334117889404297 sec
WARNING:absl:Finished tracing + transforming fn for jit in 0.0003993511199951172 sec
WARNING:absl:Compiling fn (139703279463296 for args (ShapedArray(float32[2,5]), ShapedArray(int32[], weak_type=True)).
WARNING:absl:Finished XLA compilation of fn in 0.04700160026550293 sec
WARNING:absl:Finished tracing + transforming f for jit in 0.0010411739349365234 sec
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.00015473365783691406 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209762752 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.04526233673095703 sec
WARNING:absl:Finished tracing + transforming prim_fun for jit in 0.00016546249389648438 sec
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.00014591217041015625 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209798976 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.00750732421875 sec
WARNING:absl:Finished tracing + transforming backward_pass for jit in 0.0011491775512695312 sec
WARNING:absl:Compiling backward_pass (139703209802560 for args (ShapedArray(float32[]),).
WARNING:absl:Finished XLA compilation of transpose(jvp(f)) in 0.041948556900024414 sec
(DeviceArray(55., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.00014543533325195312 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209800384 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.007508516311645508 sec
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.0001461505889892578 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209863232 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.007668972015380859 sec
WARNING:absl:Finished tracing + transforming backward_pass for jit in 0.0005974769592285156 sec
WARNING:absl:Compiling backward_pass (139703209362624 for args (ShapedArray(float32[]),).
WARNING:absl:Finished XLA compilation of transpose(jvp(f)) in 0.005425214767456055 sec
(DeviceArray(65., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.00014638900756835938 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209350720 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.007513523101806641 sec
WARNING:absl:Finished tracing + transforming <unnamed wrapped function> for jit in 0.00015473365783691406 sec
WARNING:absl:Compiling <unnamed wrapped function> (139703209372160 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.007587909698486328 sec
WARNING:absl:Finished tracing + transforming backward_pass for jit in 0.0005350112915039062 sec
WARNING:absl:Compiling backward_pass (139703209370048 for args (ShapedArray(float32[]),).
WARNING:absl:Finished XLA compilation of transpose(jvp(f)) in 0.0054433345794677734 sec
(DeviceArray(75., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
```

Comment out the checkpoint decorator leads to the wanted behavior:
```
WARNING:absl:Finished tracing + transforming prim_fun for jit in 0.0002498626708984375 sec
WARNING:absl:Finished tracing + transforming fn for jit in 0.00040721893310546875 sec
WARNING:absl:Compiling fn (140693235752000 for args (ShapedArray(float32[2,5]), ShapedArray(int32[], weak_type=True)).
WARNING:absl:Finished XLA compilation of fn in 0.04748940467834473 sec
WARNING:absl:Finished tracing + transforming f for jit in 0.0010097026824951172 sec
WARNING:absl:Compiling f (140692730754112 for args (ShapedArray(float32[2,5]),).
WARNING:absl:Finished XLA compilation of jvp(f) in 0.04457998275756836 sec
WARNING:absl:Finished tracing + transforming prim_fun for jit in 0.0001583099365234375 sec
WARNING:absl:Finished tracing + transforming backward_pass for jit in 0.0004944801330566406 sec
WARNING:absl:Compiling backward_pass (140692730730304 for args (ShapedArray(float32[]),).
WARNING:absl:Finished XLA compilation of transpose(jvp(f)) in 0.041858673095703125 sec
(DeviceArray(55., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
(DeviceArray(65., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
(DeviceArray(75., dtype=float32), DeviceArray([[1., 1., 1., 1., 1.],
             [1., 1., 1., 1., 1.]], dtype=float32))
```




Discussion:
Hey, since this is a big blocker for our project at the moment, I would be very thankful if some of the jax-internals experts could help with this issue? I tried to dig into the code why this happening but found nothing yet :(
Thanks for pinging this and highlighting that it's a blocker. Sorry for not getting to it sooner!

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -28,7 +28,7 @@
 from jax._src.ad_util import (add_jaxvals, add_jaxvals_p, zeros_like_jaxval,
                               zeros_like_aval, zeros_like_p, Zero)
 from jax._src.util import (unzip2, safe_map, safe_zip, split_list, wrap_name,
-                           as_hashable_function, cache)
+                           as_hashable_function, weakref_lru_cache)
 from jax.tree_util import register_pytree_node
 from jax import linear_util as lu
 from jax._src.api_util import flatten_fun, flatten_fun_nokwargs
@@ -586,7 +586,8 @@ def traceable(num_primals, in_tree_def, *primals_and_tangents):
 
 def call_transpose(primitive, params, call_jaxpr, args, ct, _, reduce_axes):
   all_args, in_tree_def = tree_flatten(((), args, ct))  # empty consts
-  fun = lu.hashable_partial(lu.wrap_init(backward_pass), call_jaxpr, reduce_axes, False)
+  fun = lu.hashable_partial(lu.wrap_init(backward_pass), call_jaxpr,
+                            reduce_axes, False)
   fun, out_tree = flatten_fun_nokwargs(fun, in_tree_def)
   if config.jax_experimental_name_stack:
     new_params = params
@@ -603,31 +604,29 @@ def call_transpose(primitive, params, call_jaxpr, args, ct, _, reduce_axes):
 
 def remat_transpose(params, call_jaxpr, primals_in, cotangents_in,
                     cotangent_in_avals, reduce_axes):
-  # backward_pass can only transpose linear computations, but the call_jaxpr embedded in
-  # remat contains primal (non-linear) equations too. Hence, we have to eliminate those
-  # (in this case via partial_eval) before we call into backward_pass again.
-  typed_call_jaxpr = core.ClosedJaxpr(call_jaxpr, [])
+  call_jaxpr = _close_jaxpr(call_jaxpr)
   unknowns = map(is_undefined_primal, primals_in)
-  primal_jaxpr, tangent_jaxpr, out_unknowns = \
-    pe.partial_eval_jaxpr(typed_call_jaxpr, unknowns=unknowns, instantiate=True)  # type: ignore
-
-  def do_transpose(primals_in, cotangents_in):
-    # NOTE: This is passing in undefined primals in place of tangent arguments, but it
-    #       should all work out, because we're only computing the primal part here.
-    residuals = core.jaxpr_as_fun(primal_jaxpr)(*primals_in)[len(cotangents_in):]
-    # Now that we have a purely linear jaxpr, we can transpose it
-    cotangents_out = backward_pass(
-        tangent_jaxpr.jaxpr, reduce_axes, False, (), primals_in + residuals, cotangents_in)
-    # backward_pass will return cotangents computed for all invars, but some of them
-    # are residuals appended by partial eval, so we need to skip those before we return.
-    return cotangents_out[:len(primals_in)]
-
-  flat_args, in_tree_def = tree_flatten((primals_in, cotangents_in))
-  flat_do_transpose, out_tree = flatten_fun_nokwargs(lu.wrap_init(do_transpose), in_tree_def)
-  flat_cotangents_out = pe.remat_call_p.bind(flat_do_transpose, *flat_args, **params)
+  primal_jaxpr, tangent_jaxpr, _ = \
+    pe.partial_eval_jaxpr(call_jaxpr, unknowns=unknowns, instantiate=True)  # type: ignore
+  args, in_tree_def = tree_flatten((primals_in, cotangents_in))
+  transpose = lu.hashable_partial(lu.wrap_init(_remat_transpose), primal_jaxpr,
+                                  tangent_jaxpr, reduce_axes)
+  flat_transpose, out_tree = flatten_fun_nokwargs(transpose, in_tree_def)
+  flat_cotangents_out = pe.remat_call_p.bind(flat_transpose, *args, **params)
   return tree_unflatten(out_tree(), flat_cotangents_out)
 primitive_transposes[pe.remat_call_p] = remat_transpose
 
+def _remat_transpose(primal_jaxpr, tangent_jaxpr, reduce_axes,
+                     primals_in, cotangents_in):
+  res = core.jaxpr_as_fun(primal_jaxpr)(*primals_in)[len(cotangents_in):]
+  cotangents_out = backward_pass(tangent_jaxpr.jaxpr, reduce_axes, False, (),
+                                 (*primals_in, *res), cotangents_in)
+  return cotangents_out[:len(primals_in)]
+
+@weakref_lru_cache
+def _close_jaxpr(jaxpr: core.Jaxpr) -> core.ClosedJaxpr:
+  return core.ClosedJaxpr(jaxpr, [])
+
 @lu.transformation_with_aux
 def nonzero_outputs(*args, **kwargs):
   results = yield args, kwargs
@@ -680,7 +679,7 @@ def jvp_jaxpr(jaxpr, nonzeros, instantiate):
   inst = tuple(instantiate) if isinstance(instantiate, list) else instantiate
   return _jvp_jaxpr(jaxpr, tuple(nonzeros), inst)
 
-@cache()
+@weakref_lru_cache
 def _jvp_jaxpr(jaxpr, nonzeros, instantiate):
   assert len(jaxpr.in_avals) == len(nonzeros)
   f = lu.wrap_init(core.jaxpr_as_fun(jaxpr))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4328,10 +4328,19 @@ def test_linearize_caching(self):
     identity = jax.checkpoint(jax.jit(lambda x: 2 * x))
     _, f_lin = jax.linearize(identity, 1.)
     with jtu.count_jit_and_pmap_compiles() as count:  # noqa: F841
-      for _ in range(10):
+      for _ in range(20):
         f_lin(1.).block_until_ready()
     self.assertEqual(count[0], 1)  # cached after first execution
 
+  def test_vjp_caching(self):
+    # https://github.com/google/jax/issues/9661
+    identity = jax.checkpoint(jax.jit(lambda x: 2 * x))
+    _, f_vjp = jax.vjp(identity, 1.)
+    with jtu.count_jit_and_pmap_compiles() as count:  # noqa: F841
+      for _ in range(20):
+        f_vjp(1.)[0].block_until_ready()
+    self.assertEqual(count[0], 2)  # eval_jaxpr on fwd, backward_pass on bwd
+
 
 class JaxprTest(jtu.JaxTestCase):