Repository: google/jax
Issue: Flaky test failure of RematTest in api_test

Run, with jaxlib built from head:
```
bazel test --//jax:build_jaxlib=false //tests:api_test_cpu --test_arg=RematTest.test_remat_grad_python_control_flow_static_argnums --test_sharding_strategy=disabled --notrim_test_configuration --runs_per_test=100
```

```
//tests:api_test_cpu                                                     FAILED in 2 out of 100 in 7.4s
  Stats over 100 runs: max = 7.4s, min = 2.0s, avg = 6.8s, dev = 1.0s
  /usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/testlogs/tests/api_test_cpu/run_17_of_100/test.log
  /usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/testlogs/tests/api_test_cpu/run_91_of_100/test.log

```

Example failure:

```
[  FAILED  ] RematTest.test_remat_grad_python_control_flow_static_argnums
======================================================================
FAIL: test_remat_grad_python_control_flow_static_argnums (__main__.RematTest)
RematTest.test_remat_grad_python_control_flow_static_argnums
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/bin/tests/api_test_cpu.runfiles/__main__/tests/api_test.py", line 4037, in test_remat_grad_python_control_flow_static_argnums
    self.assertAllClose(ans, expected, check_dtypes=False)
  File "/usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/bin/tests/api_test_cpu.runfiles/__main__/jax/_src/test_util.py", line 798, in assertAllClose
    self.assertArraysAllClose(x, y, check_dtypes=False, atol=atol, rtol=rtol,
  File "/usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/bin/tests/api_test_cpu.runfiles/__main__/jax/_src/test_util.py", line 763, in assertArraysAllClose
    _assert_numpy_allclose(x, y, atol=atol, rtol=rtol, err_msg=err_msg)
  File "/usr/local/google/home/phawkins/.cache/bazel/_bazel_phawkins/f5f850d43f4a41d5e22a7e7660064506/execroot/__main__/bazel-out/k8-opt/bin/tests/api_test_cpu.runfiles/__main__/jax/_src/public_test_util.py", line 96, in _assert_numpy_allclose
    np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
  File "/usr/local/google/home/phawkins/.pyenv/versions/py3913/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 1527, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/usr/local/google/home/phawkins/.pyenv/versions/py3913/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-06, atol=1e-06

Mismatched elements: 1 / 1 (100%)
Max absolute difference: 0.41614684
Max relative difference: 1.
 x: array(0., dtype=float32)
 y: array(-0.416147)

----------------------------------------------------------------------
Ran 1 test in 0.352s

FAILED (failures=1)
```


It doesn't appear to repro with jaxlib from pypi.

Now, I'm not convinced this is an XLA bug from the evidence. Interestingly we dump different XLA computations with `XLA_FLAGS=--xla_dump_to` in cases when we fail and cases when we succeed! So that suggests to me this is a JAX bug.

I've seen this on CPU, GPU and, TPU.


Changed files:
jax/_src/ad_checkpoint.py

diff --git a/jax/_src/ad_checkpoint.py b/jax/_src/ad_checkpoint.py
index 7af87c12f..9ce4b8a22 100644
--- a/jax/_src/ad_checkpoint.py
+++ b/jax/_src/ad_checkpoint.py
@@ -14,7 +14,8 @@
 
 from functools import partial
 import operator as op
-from typing import Callable, Optional, List, Tuple, Sequence, Set, Union, Any
+from typing import (Callable, Optional, List, Tuple, Sequence, Set, Union, Any,
+                    FrozenSet)
 import types
 
 from absl import logging
@@ -302,7 +303,7 @@ def _remat_static_argnums(fun, static_argnums, args):
 
 class WrapHashably:
   val: Any
-  hash: Optional[int] = None
+  hash: int
   hashable: bool
 
   def __init__(self, val):
@@ -317,8 +318,10 @@ class WrapHashably:
     return self.hash
   def __eq__(self, other):
     if isinstance(other, WrapHashably):
-      try: return self.val == other.val
-      except: return self.val is other.val
+      if self.hashable and other.hashable:
+        return self.val == other.val
+      else:
+        return self.val is other.val
     return False
 
 # This caching is useful to avoid retracing even when static_argnums is used.
@@ -326,7 +329,7 @@ class WrapHashably:
 # On that benchmark, including this caching makes a ~10x difference (which can
 # be made arbitrary large by involving larger functions to be traced).
 @weakref_lru_cache
-def _dyn_args_fun(fun: Callable, static_argnums: Tuple[int, ...],
+def _dyn_args_fun(fun: Callable, static_argnums: FrozenSet[int],
                   static_args: Tuple[WrapHashably, ...], nargs: int):
   def new_fun(*dyn_args, **kwargs):
     static_args_, dyn_args_ = iter(static_args), iter(dyn_args)