Repository: google/jax

Issue:
Memory leak when computing the hessian of a jitted function
Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

```python
import jax
import jax.numpy as jnp

f = lambda x: jnp.square(x).mean()
jf = jax.jit(f)

x = jax.random.uniform(jax.random.PRNGKey(0), shape=(8, 4))

# Start monitoring memory with htop
while True:
    try:
        # let this run for a bit, notice memory usage doesn't increase
        y = jax.hessian(f)(x)
    except KeyboardInterrupt:
        pass

while True:
    # let this run, notice memory usage starts increasing immediately
    # also, the memory usage doesn't appear to stop increasing ever
    y = jax.hessian(jf)(x)
```
Also, note that the equivalent code with `jax.grad` (or `jax.jacfwd` for multidimensional functions) replacing `jax.hessian` *does not* cause a memory leak.

As a result of this leak, my python process ends up getting killed by my OS (due to OOM), so I don't ever see any python or jax traces.


Discussion:
Thanks for ~the question~ raising this, and with a clear reproducer!

At first I thought could be an artifact of [asynchronous dispatch](https://jax.readthedocs.io/en/latest/async_dispatch.html). (The `jit` dispatch path is different from the non-`jit` dispatch path, so that may play a role in why `f` and `jf` behave differently.)

But I ran this version, which forces synchronizations, and it seems like memory is still growing without bound (on the CPU backend):

```python
while True:
    jax.hessian(jf)(x).block_until_ready()
```

I then tried running the script like this:

```
env JAX_LOG_COMPILES=1 python issue7621.py
```

and saw that we're doing lots of recompiles here:

```
...
WARNING:absl:Compiling backward_pass (140099626713344) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626782592) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626820800) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626782144) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626862272) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626817664) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626755712) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626837568) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626919552) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626921344) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626884480) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626944512) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626883712) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
WARNING:absl:Compiling backward_pass (140099626911936) for args (ShapedArray(float32[8,4]), ShapedArray(float32[]), ShapedArray(float32[1]), ShapedArray(float32[32,8,4])).
...
```

Maybe that should've been obvious from your clue about `grad`/`jacfwd` behaving differently.

Somehow `hessian` is causing compilation cache misses here...
@mattjj Nice catch, I wasn't aware of the `JAX_LOG_COMPILES` variable. I wonder if declaring the function `h = hessian(jf)` outside the loop and calling `h(x)` within the loop would still cause the issue. I'll give it a try (though admittedly I don't know much about how `jit` works, so this may not make sense ;) )

**Edit**: the hessian outside the loop trick doesn't work, though that probably isn't surprising.
This _should_ be getting cache hits, so it's definitely a bug. The situation should be basically the same as for `grad`, so there's some difference between `grad` and `hessian` which we need to pin down...
I see the issue. It's a problem with `vmap`. It's a bit hard to unpack in full detail, so for now I'm just going to leave some notes for myself and other jax devs: in [this cache key](https://github.com/google/jax/blob/a3387aeb0c26e8334d874160a00db7d5e174264b/jax/linear_util.py#L255), on each invocation we're getting different values because `key[0][7][1][0].payload` is a different [temporary axis name](https://github.com/google/jax/blob/a3387aeb0c26e8334d874160a00db7d5e174264b/jax/_src/api.py#L1259) each time. That's happening in turn because `jacrev's` [`pullback`](https://github.com/google/jax/blob/a3387aeb0c26e8334d874160a00db7d5e174264b/jax/_src/api.py#L1033-L1035), which gets vmapped, is a fresh Python callable each time.
More notes to self:

We added temporary axis names in #917 for soft_pmap (now replaced by xmap). But the issue here is about vmap using temporary axis names.

The reason we needed them for soft_pmap/xmap is that those implementations split axes, and so they need to generate new names which could be used in collectives (and hence have to be distinguished from one another, etc).

But here with vmap, when we aren't given a name, the axis can't be involved in any collectives. So it seems like we should just leave it as None. Unless I'm forgetting something...
Adam and I added this TempAxisName stuff to vmap in #4005. I wonder if that was just a bug...

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1256,8 +1256,6 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
     docstr += "\n\nOriginal documentation:\n\n"
     docstr += fun.__doc__
 
-  axis_name = core._TempAxisName(fun) if axis_name is None else axis_name
-
   if isinstance(in_axes, list):
     # To be a tree prefix of the positional args tuple, in_axes can never be a
     # list: if in_axes is not a leaf, it must be a tuple of trees. However,
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -66,7 +66,7 @@ def _match_axes(axis_size, axis_name, in_dims, out_dims_thunk, out_dim_dests,
   out_dims = out_dims_thunk()
   for od, od_dest in zip(out_dims, out_dim_dests):
     if od is not None and not isinstance(od_dest, int):
-      if not isinstance(axis_name, core._TempAxisName):
+      if not isinstance(axis_name, core._TempAxisName) and axis_name is not None:
         msg = f"vmap has mapped output (axis_name={axis_name}) but out_axes is {od_dest}"
       else:
         msg = f"vmap has mapped output but out_axes is {od_dest}"


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2949,6 +2949,20 @@ def loss(A, x):
     with jax.checking_leaks():
       _ = jax.grad(loss)(A, x)  # doesn't crash
 
+  def test_vmap_caching(self):
+    # https://github.com/google/jax/issues/7621
+
+    f = lambda x: jnp.square(x).mean()
+    jf = jax.jit(f)
+    x = jax.random.uniform(jax.random.PRNGKey(0), shape=(8, 4))
+
+    with jtu.count_jit_and_pmap_compiles() as count:  # noqa: F841
+      jax.hessian(jf)(x).block_until_ready()
+      jax.hessian(jf)(x).block_until_ready()
+      jax.hessian(jf)(x).block_until_ready()
+
+    self.assertEqual(count[0], 2)
+
 
 class RematTest(jtu.JaxTestCase):