Repository: google/jax

Issue:
Store empty exception when rerunning a computation that failed a NaN check
```
import jax
try:
 with jax.debug_nans(True):
   jax.grad(jax.numpy.linalg.norm)(jax.numpy.zeros((3, 3)))
except FloatingPointError as e:
  pass
with jax.debug_nans(True):
 jax.grad(jax.numpy.linalg.norm)(jax.numpy.zeros((3, 3)))
```

```
Traceback (most recent call last):
  File "/Users/phawkins/p/jax/tt.py", line 8, in <module>
    jax.grad(jax.numpy.linalg.norm)(jax.numpy.zeros((3, 3)))
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 140, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 810, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 140, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 879, in value_and_grad_f
    g = vjp_py(np.ones((), dtype=dtype))
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 1872, in _vjp_pullback_wrapper
    ans = fun(*args)
  File "/Users/phawkins/p/jax/jax/interpreters/ad.py", line 121, in unbound_vjp
    arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
  File "/Users/phawkins/p/jax/jax/interpreters/ad.py", line 215, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(
  File "/Users/phawkins/p/jax/jax/interpreters/ad.py", line 542, in call_transpose
    out_flat = primitive.bind(fun, *all_args, **new_params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1561, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1552, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1564, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 605, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/Users/phawkins/p/jax/jax/interpreters/xla.py", line 577, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 262, in memoized_fun
    fun.populate_stores(stores)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 150, in populate_stores
    self_store.store(other_store.val)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 106, in val
    raise StoreException("Store empty")
jax.linear_util.StoreException: Store empty
```


Discussion:
Simplified repro:

```python
import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  return jnp.divide(x, x)

try:
 with jax.debug_nans(True):
   jax.grad(f)(0.)
except FloatingPointError as e:
  pass
with jax.debug_nans(True):
  jax.grad(f)(0.)
```
We're corrupting the cache by [clearing the stores here](https://github.com/google/jax/blob/712d5f7e2bc99d2a1856a453f190da69d855b98b/jax/interpreters/xla.py#L593). I am dumb for not realizing this earlier.

We could remove the cache entry when we clear the stores. Or we could find some other way to suppress the 'store occupied' errors that this line was originally meant to avoid.

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -619,7 +619,7 @@ def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name,
   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
                                *unsafe_map(arg_spec, args))
   try:
-    return compiled_fun(*args)
+    out = compiled_fun(*args)
   except FloatingPointError:
     assert config.jax_debug_nans or config.jax_debug_infs  # compiled_fun can only raise in this case
     print("Invalid value encountered in the output of a jit function. "
@@ -630,11 +630,13 @@ def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name,
     # but which config.jax_debug_nans is meant to opt into, we'll be re-executing
     # any linear_util.py-style side effects, i.e. re-populating Stores created
     # by any transformation_with_aux's applied to fun. Since this is
-    # intentional here, to avoid "Store occupied" errors we reset the stores to
-    # be empty.
-    for store in fun.stores: store and store.reset()
+    # intentional here, to avoid "Store occupied" errors we clone the WrappedFun
+    # with empty stores.
+    stores = [lu.Store() for _ in fun.stores]
+    clone = lu.WrappedFun(fun.f, fun.transforms, stores, fun.params)
     with core.new_sublevel():
-      return fun.call_wrapped(*args)  # probably won't return
+      _ = clone.call_wrapped(*args)  # probably won't return
+  return out
 
 def flatten_shape(s: XlaShape) -> Sequence[Tuple[Sequence[int], XlaShape]]:
   """Expands a given shape tree into a flat list of indices to arrays.


Test patch:
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
--- a/tests/debug_nans_test.py
+++ b/tests/debug_nans_test.py
@@ -201,5 +201,18 @@ def f(x):
       with self.assertRaisesRegex(FloatingPointError, msg):
         f(1)
 
+  def testDebugNansDoesntCorruptCaches(self):
+    # https://github.com/google/jax/issues/6614
+    @jax.jit
+    def f(x):
+      return jnp.divide(x, x)
+
+    for _ in range(2):
+      try:
+       with jax.debug_nans(True):
+         jax.grad(f)(0.)
+      except FloatingPointError:
+        pass
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())