Repository: google/jax
Issue: `jax.nn.sigmoid` custom jvp leak

Hello,

I got a tracing leak related to `jax.nn.sigmoid` custom jvp. I think this is related to https://github.com/google/jax/issues/5636.

I'm using JAX 0.2.21 on CPU.

```python
import jax
import jax.numpy as jnp
print(jax.__version__)
# 0.2.21

def _sigmoid(x: jnp.ndarray):
    return 1.0 / (1. + jnp.exp(-x))


@jax.jit
def fwd():
    a = jnp.array(1.)

    def f(hx, _):
        hx = jax.nn.sigmoid(hx + a)
        # hx = _sigmoid(hx + a)
        return hx, None

    hx = jnp.array(0.)
    jax.lax.scan(f, hx, None, length=2)

with jax.checking_leaks():
    fwd()

# [...]
# Exception: Leaked sublevel 1. Leaked tracer(s): [Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>].
```
When I replace `jax.nn.sigmoid` by `_sigmoid`, no leak detected.

The full stack trace:

```
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-22-61b72614d9a4> in <module>
     16 with jax.checking_leaks():
---> 17     fwd()

    [... skipping hidden 1 frame]

~/.local/lib/python3.9/site-packages/jax/_src/api.py in cache_miss(*args, **kwargs)
    410     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 411     out_flat = xla.xla_call(
    412         flat_fun, *args_flat,

~/.local/lib/python3.9/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1617   def bind(self, fun, *args, **params):
-> 1618     return call_bind(self, fun, *args, **params)
   1619 

~/.local/lib/python3.9/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1608   tracers = map(top_trace.full_raise, args)
-> 1609   outs = primitive.process(top_trace, fun, tracers, params)
   1610   return map(full_lower, apply_todos(env_trace_todo(), outs))

~/.local/lib/python3.9/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1620   def process(self, trace, fun, tracers, params):
-> 1621     return trace.process_call(self, fun, tracers, params)
   1622 

~/.local/lib/python3.9/site-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
    614   def process_call(self, primitive, f, tracers, params):
--> 615     return primitive.impl(f, *tracers, **params)
    616   process_map = process_call

~/.local/lib/python3.9/site-packages/jax/interpreters/xla.py in _xla_call_impl(***failed resolving arguments***)
    621   del inline  # Only used at tracing time
--> 622   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
    623                                *unsafe_map(arg_spec, args))

~/.local/lib/python3.9/site-packages/jax/linear_util.py in memoized_fun(fun, *args)
    261     else:
--> 262       ans = call(fun, *args)
    263       cache[key] = (ans, fun.stores)

~/.local/lib/python3.9/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    693 def _xla_callable(fun: lu.WrappedFun, device, backend, name, donated_invars, *arg_specs):
--> 694   return lower_xla_callable(fun, device, backend, name, donated_invars, *arg_specs).compile().unsafe_call
    695 

~/.local/lib/python3.9/site-packages/jax/interpreters/xla.py in lower_xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    701   abstract_args, arg_devices = unzip2(arg_specs)
--> 702   jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
    703       fun, abstract_args, pe.debug_info_final(fun, "jit"))

~/.local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals, debug_info)
   1521     with core.new_sublevel():
-> 1522       jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1523     del fun, main

/usr/lib/python3.9/contextlib.py in __exit__(self, type, value, traceback)
    123             try:
--> 124                 next(self.gen)
    125             except StopIteration:

~/.local/lib/python3.9/site-packages/jax/core.py in new_sublevel()
    842       if leaked_tracers:
--> 843         raise Exception(f'Leaked sublevel {t()}. Leaked tracer(s): {leaked_tracers}.')
    844 

UnfilteredStackTrace: Exception: Leaked sublevel 1. Leaked tracer(s): [Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>].

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Exception                                 Traceback (most recent call last)
<ipython-input-22-61b72614d9a4> in <module>
     15 
     16 with jax.checking_leaks():
---> 17     fwd()

/usr/lib/python3.9/contextlib.py in __exit__(self, type, value, traceback)
    122         if type is None:
    123             try:
--> 124                 next(self.gen)
    125             except StopIteration:
    126                 return False

Exception: Leaked sublevel 1. Leaked tracer(s): [Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>].
```


Changed files:
jax/interpreters/partial_eval.py, tests/nn_test.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 2104fa183..93e62fe4d 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -1408,8 +1408,9 @@ class DynamicJaxprTrace(core.Trace):
     with core.new_sublevel():
       fun_jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, self.main, in_avals)
     closed_fun_jaxpr = core.ClosedJaxpr(convert_constvars_jaxpr(fun_jaxpr), ())
+    main_ = ref(self.main)
     jvp_jaxpr_thunk = _memoize(
-        lambda: trace_to_subjaxpr_dynamic(jvp, self.main, 2 * in_avals)[::2])
+        lambda: trace_to_subjaxpr_dynamic(jvp, main_(), 2 * in_avals)[::2])
     out_tracers = [DynamicJaxprTracer(self, a) for a in out_avals]
     invars = map(self.getvar, tracers)
     constvars = map(self.getvar, map(self.instantiate_const, consts))
@@ -1430,8 +1431,9 @@ class DynamicJaxprTrace(core.Trace):
     with core.new_sublevel():
       fun_jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, self.main, in_avals)
     closed_fun_jaxpr = core.ClosedJaxpr(convert_constvars_jaxpr(fun_jaxpr), ())
+    main_ = ref(self.main)
     fwd_jaxpr_thunk = _memoize(
-        lambda: trace_to_subjaxpr_dynamic(fwd, self.main, in_avals)[::2])
+        lambda: trace_to_subjaxpr_dynamic(fwd, main_(), in_avals)[::2])
     out_tracers = [DynamicJaxprTracer(self, a) for a in out_avals]
     invars = map(self.getvar, tracers)
     constvars = map(self.getvar, map(self.instantiate_const, consts))
@@ -1449,6 +1451,9 @@ class DynamicJaxprTrace(core.Trace):
     assert False  # unreachable
 
 def _memoize(thunk):
+  if config.jax_check_tracer_leaks:
+    return thunk
+
   cell = []
   saved_state = core.thread_local_state.trace_state.copy()
   def memoized():
diff --git a/tests/nn_test.py b/tests/nn_test.py
index 12861dd05..b29a7edc3 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -176,6 +176,22 @@ class NNFunctionsTest(jtu.JaxTestCase):
   def testTanhExists(self):
     nn.tanh  # doesn't crash
 
+  def testCustomJVPLeak(self):
+    # https://github.com/google/jax/issues/8171
+    @jax.jit
+    def fwd():
+      a = jnp.array(1.)
+
+      def f(hx, _):
+        hx = jax.nn.sigmoid(hx + a)
+        return hx, None
+
+      hx = jnp.array(0.)
+      jax.lax.scan(f, hx, None, length=2)
+
+    with jax.checking_leaks():
+      fwd()  # doesn't crash
+
 InitializerRecord = collections.namedtuple(
   "InitializerRecord",
   ["name", "initializer", "shapes", "dtypes"])