Repository: google/jax

Issue:
Spurious leak checker error with custom JVP
I'd love to be able to use the (great-looking) new leaked tracer tools, but I'm running into a problem. After I set
```
export JAX_CHECK_TRACER_LEAKS=1
```
and then run:
```python
import jax

@jax.custom_jvp
def t(y):
    return y

def t_jvp(primals, tangents):
    assert False

t.defjvp(t_jvp)

@jax.jit
def s(y):
    return t(y)

s(0.0)
```
I unfortunately get
```
Exception: Leaked trace MainTrace(0,DynamicJaxprTrace)
```
There's possibly a deeper bug here because I'm getting leaked tracers in my code wtihout the flag, so I'm hoping this is a spurious error.


Discussion:
Thanks for this repro! The leak checker definitely tends towards false positives right now, so it's a good a priori guess.
Tricky! There are at least two issues, both pertaining to [the same line](https://github.com/google/jax/blob/0f7cc80d38822c8753e076ff6d12022d969c2dbb/jax/interpreters/partial_eval.py#L1131-L1132):
1. That `_memoize` decorator is acting like a cache, and it's a cache we forgot to disable along with the others in #5492 when the leak checker is active. It may be a simple matter of disabling that memoization; I don't think it's necessary for correctness, though I have paged this out and could be wrong about that. All the tests pass, at least!
2. That [`lambda`](https://github.com/google/jax/blob/0f7cc80d38822c8753e076ff6d12022d969c2dbb/jax/interpreters/partial_eval.py#L1132) being created on that line keeps a reference to the MainTrace instance, and it's not being cleaned up fast enough, so the leak checker spots it.

I think both may only arise in the case where the jvp rule raises an exception.

The fix to (1) is easy, as with the other caches:

```python
def _memoize(thunk):
  if core.debug_state.check_leaks:
    return thunk

  cell = []
  saved_state = core.thread_local_state.trace_state.copy()
  def memoized():
    nonlocal thunk
    if not cell:
      prev_state = core.thread_local_state.trace_state
      core.thread_local_state.trace_state = saved_state
      try:
        cell.append(thunk())
      finally:
        del thunk
        core.thread_local_state.trace_state = prev_state
    return cell[0]
  return memoized
```

 The fix to (2) I'm not sure about yet...
(I can't work on this more at the moment but wanted to spend a few mins on an initial investigation!)
Thanks for looking into this!  Just so you know, I did find the problem (I was jusing `jax.jit` instead of `hk.jit`), so this isn't holding me up anymore.  Hope the MWE is useful to you though :)
Woo, that's great to hear! Sounds like a common issue as well. We're working with the Haiku folks to improve things.

cc @LenaMartens just to bring this leak checker thread to her attention
@mattjj That's great to hear!  Just to keep you informed, there's a beautiful solution to my problem by shoyer here https://github.com/google/jax/pull/4117.  I'm just now trying to code up the Haiku side of things to get it to work locally.

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -18,6 +18,7 @@
 from contextlib import contextmanager
 from collections import namedtuple
 from functools import total_ordering
+import gc
 import itertools as it
 from weakref import ref
 import threading
@@ -736,6 +737,17 @@ def reset_trace_state() -> bool:
 def cur_sublevel() -> Sublevel:
   return thread_local_state.trace_state.substack[-1]
 
+def maybe_find_leaked_tracers(x: Optional[Union[MainTrace, Sublevel]]):
+  """Find the leaked tracers holding a reference to the MainTrace or SubLevel.
+
+  It's possible there's none! eg. there's some cases where JAX itself holds a
+  reference to `x` inside of a lambda closure, and no tracers were leaked
+  by the user. In this case an empty list is returned.
+  """
+  traces = list(filter(lambda x: isinstance(x, Trace), gc.get_referrers(x)))
+  tracers = list(filter(lambda x: isinstance(x, Tracer), gc.get_referrers(*traces)))
+  return tracers
+
 @contextmanager
 def new_main(trace_type: Type[Trace],
              dynamic: bool = False,
@@ -761,7 +773,9 @@ def new_main(trace_type: Type[Trace],
     t = ref(main)
     del main
     if t() is not None:
-      raise Exception(f'Leaked trace {t()}')
+      leaked_tracers = maybe_find_leaked_tracers(t())
+      if leaked_tracers:
+        raise Exception(f'Leaked level {t()}. Leaked tracer(s): {leaked_tracers}.')
 
 @contextmanager
 def new_base_main(trace_type: Type[Trace]) -> Generator[MainTrace, None, None]:
@@ -782,7 +796,9 @@ def new_base_main(trace_type: Type[Trace]) -> Generator[MainTrace, None, None]:
     t = ref(main)
     del main
     if t() is not None:
-      raise Exception('Leaked trace {}'.format(t()))
+      leaked_tracers = maybe_find_leaked_tracers(t())
+      if leaked_tracers:
+        raise Exception(f'Leaked level {t()}. Leaked tracer(s): {leaked_tracers}.')
 
 @contextmanager
 def eval_context():
@@ -802,7 +818,9 @@ def new_sublevel() -> Generator[None, None, None]:
     t = ref(sublevel)
     del sublevel
     if t() is not None:
-      raise Exception(f'Leaked sublevel {t()}.')
+      leaked_tracers = maybe_find_leaked_tracers(t())
+      if leaked_tracers:
+        raise Exception(f'Leaked sublevel {t()}. Leaked tracer(s): {leaked_tracers}.')
 
 def full_lower(val):
   if isinstance(val, Tracer):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2596,6 +2596,23 @@ def g(x):
       with self.assertRaisesRegex(Exception, r"Leaked sublevel"):
         f(3)
 
+  def test_leak_checker_avoids_false_positive_custom_jvp(self):
+    # see https://github.com/google/jax/issues/5636
+    with jax.checking_leaks():
+      @api.custom_jvp
+      def t(y):
+        return y
+
+      def t_jvp(p, t):
+        pass
+
+      t.defjvp(t_jvp)
+
+      @jit
+      def s(y):
+        return t(y)
+      s(3)  # doesn't crash
+
   def test_default_backend(self):
     first_local_device = api.local_devices()[0]
     self.assertEqual(first_local_device.platform, api.default_backend())