Repository: google/jax

Issue:
experimental.enable_x64 fails for random.uniform
Short repro:
```python
from jax import random
from jax.experimental import disable_x64, enable_x64

f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
with disable_x64():
    f()
with enable_x64():
    f()
    f()  # <--- fails
```
Failure:
```pytb
Traceback (most recent call last):
  File "tmp.py", line 11, in <module>
    f()  # <--- fails
  File "tmp.py", line 4, in <lambda>
    f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
  File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
    return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "tmp.py", line 11, in <module>
    f()  # <--- fails
  File "tmp.py", line 4, in <lambda>
    f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
  File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
    return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
  File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
```


Discussion:
This is not a problem with the python JIT path, but is a problem with the CPP jit. I think the issue is here: https://github.com/google/jax/blob/f4b5ff9d466c3f93d850efd772eee0ef03601a05/jax/api.py#L370-L372

The cpp JIT determines the X64 flag value only the first time the CPP function is called: https://github.com/tensorflow/tensorflow/blob/26bb688b0cba1e56fdee7ecde7ebde0d72ce2066/tensorflow/compiler/xla/python/jax_jit.cc#L1108
and the result is stored in an `absl` constant, which I believe essentially acts as a global constant: https://github.com/tensorflow/tensorflow/blob/26bb688b0cba1e56fdee7ecde7ebde0d72ce2066/tensorflow/compiler/xla/python/jax_jit.cc#L900

This means that compiling the function twice with different X64 flags can make the second compiled function fail on subsequent calls because the global cache of the X64 flag reflects the value at the first rather than the second compilation.

I think that fixing this will require changing the C++ side to relax the assumption that `jax_enable_x64` is constant within a session.
Absolutely.

Just to check on the semantics Is the context manager expected to change the code at run time or definition time? (I strongly expect at runtime).

```
with enablex64():
  f= jax.jit(...)
  f(x) => I expect it should be with x64?
f(x) => I imagine it should be with whatever global value is define in the program?
with disablex64():
  f(x) => I imagine it should be x32?
```

Assuming this semantics, it means the C++ will need for each call to know what is the local value for enable_x64.
There are a few options:
- Call Python each time, from C++, to get the value of the flag. I dislike this approach, as the goal is to call the least python as Possible
- Pass the value from Python to C++, e.g. as the first argument, as take it from C++, and consider the rest of the arguments as the jitted function call.
- Have jax_enable_64 live in C++, and have the context manager update this value, and access this value from C++. This is what is done for DisableJit:
https://github.com/tensorflow/tensorflow/blob/26bb688b0cba1e56fdee7ecde7ebde0d72ce2066/tensorflow/compiler/xla/python/jax_jit.cc#L1141

The difficulty of that option, is that we need to move the value from the config to C++ (i.e. when we update the flags, we should set the C++ value).
> Is the context manager expected to change the code at run time or definition time? (I strongly expect at runtime).

Indeed, the semantics are that that the code is changed at run-time. For the CPP jit, we accomplish this by using the current value of the X64 flag as part of the cache key: https://github.com/google/jax/blob/f4b5ff9d466c3f93d850efd772eee0ef03601a05/jax/_src/util.py#L187-L203

> Assuming this semantics, it means the C++ will need for each call to know what is the local value for enable_x64.

This is True. Our initial attempt at that was to add the X64 flag to the context used to define the cache key: https://github.com/google/jax/blob/f4b5ff9d466c3f93d850efd772eee0ef03601a05/jax/api.py#L378-L379
But obviously this has to be threaded through more deeply. I think the second approach is probably the best, as it basically matches the approach in the python JIT and avoids extra complication.

> - Pass the value from Python to C++, e.g. as the first argument, as take it from C++, and consider the rest of the arguments as the jitted function call.

In fact, the value is already available to C++ as part of the `context` variable as shown above... though it would make sense to make this more explicit in the function's contract if we go this route.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -355,7 +355,11 @@ def get_jax_enable_x64():
     functions decorated with jax.jit), so we delay inspecting the value
     of the jax_enable_x64 flag until JIT time.
     """
-    return config.x64_enabled
+    # TODO(jblespiau): Delete when jaxlib 0.1.62 is the minimal version.
+    if lib._xla_extension_version >= 4:
+      return config.read("jax_enable_x64")
+    else:
+      return config.x64_enabled
 
   def get_jax_disable_jit_flag():
     """Returns the value of the `jax_disable_jit` flag.
@@ -375,8 +379,10 @@ def get_jax_disable_jit_flag():
   @wraps(fun)
   @api_boundary
   def f_jitted(*args, **kwargs):
+    # TODO(jblespiau): We can remove `config.x64_enabled` when jaxlib has
+    # extension version 4
     context = (getattr(core.thread_local_state.trace_state.trace_stack,
-                       'dynamic', None), config.x64_enabled)
+                         "dynamic", None), config.x64_enabled)
     # TODO(jblespiau): Move this to C++.
     if (FLAGS.jax_debug_nans or FLAGS.jax_debug_infs) and not _jit_is_disabled():
       device_arrays = cpp_jitted_f(context, *args, **kwargs)
diff --git a/jax/config.py b/jax/config.py
--- a/jax/config.py
+++ b/jax/config.py
@@ -16,6 +16,7 @@
 import sys
 import threading
 from typing import Optional
+from jax import lib
 
 def bool_env(varname: str, default: bool) -> bool:
   """Read an environment variable and interpret it as a boolean.
@@ -49,6 +50,7 @@ def __init__(self):
 
 
 class Config:
+  # TODO(jakevdp): Remove when minimum jaxlib is has extension version 4
   _thread_local_state = _ThreadLocalState()
 
   def __init__(self):
@@ -149,13 +151,23 @@ def disable_omnistaging(self):
 
   @property
   def x64_enabled(self):
-    if self._thread_local_state.enable_x64 is None:
-      self._thread_local_state.enable_x64 = bool(self.read('jax_enable_x64'))
-    return self._thread_local_state.enable_x64
+   if lib._xla_extension_version >= 4:
+     if lib.jax_jit.get_enable_x64() is None:
+       lib.jax_jit.set_enable_x64(bool(self.read('jax_enable_x64')))
+     return lib.jax_jit.get_enable_x64()
+   else:
+     # TODO(jakevdp): Remove when minimum jaxlib is has extension version 4
+     if self._thread_local_state.enable_x64 is None:
+       self._thread_local_state.enable_x64 = bool(self.read('jax_enable_x64'))
+     return self._thread_local_state.enable_x64
 
   # TODO(jakevdp): make this public when thread-local x64 is fully implemented.
   def _set_x64_enabled(self, state):
-    self._thread_local_state.enable_x64 = bool(state)
+    if lib._xla_extension_version >= 4:
+      lib.jax_jit.set_enable_x64(bool(state))
+    else:
+      # TODO(jakevdp): Remove when minimum jaxlib is has extension version 4
+      self._thread_local_state.enable_x64 = bool(state)
 
 
 class NameSpace(object):
diff --git a/jax/experimental/x64_context.py b/jax/experimental/x64_context.py
--- a/jax/experimental/x64_context.py
+++ b/jax/experimental/x64_context.py
@@ -18,7 +18,7 @@
 """
 
 from contextlib import contextmanager
-from jax.config import config
+from jax import config
 
 @contextmanager
 def enable_x64():


Test patch:
diff --git a/tests/x64_context_test.py b/tests/x64_context_test.py
--- a/tests/x64_context_test.py
+++ b/tests/x64_context_test.py
@@ -16,10 +16,16 @@
 import concurrent.futures
 import time
 
-from absl.testing import absltest, parameterized
-
-from jax import api, lax, partial, random
-from jax.config import config, FLAGS
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import jax
+from jax import api
+from jax import lax
+from jax import partial
+from jax import random
+from jax.config import config
+from jax.config import FLAGS
 from jax.experimental import enable_x64, disable_x64
 import jax.numpy as jnp
 import jax.test_util as jtu
@@ -48,11 +54,35 @@ def test_make_array(self, jit):
     func = _maybe_jit(jit, lambda: jnp.arange(10.0))
     dtype_start = func().dtype
     with enable_x64():
-      self.assertEqual(func().dtype, 'float64')
+      self.assertEqual(func().dtype, "float64")
     with disable_x64():
-      self.assertEqual(func().dtype, 'float32')
+      self.assertEqual(func().dtype, "float32")
     self.assertEqual(func().dtype, dtype_start)
 
+  @parameterized.named_parameters(
+      jtu.cases_from_list({
+          "testcase_name": "_jit={}_f_{}".format(jit, f.__name__),
+          "jit": jit,
+          "enable_or_disable": f
+      } for jit in ["python", "cpp", None] for f in [enable_x64, disable_x64]))
+  def test_correctly_capture_default(self, jit, enable_or_disable):
+    if jit == "cpp" and not config.omnistaging_enabled:
+      self.skipTest("cpp_jit requires omnistaging")
+
+    # The fact we defined a jitted function with a block with a different value
+    # of `config.enable_x64` has no impact on the output.
+    with enable_or_disable():
+      func = _maybe_jit(jit, lambda: jnp.arange(10.0))
+      func()
+
+    expected_dtype = "float64" if config.read("jax_enable_x64") else "float32"
+    self.assertEqual(func().dtype, expected_dtype)
+
+    with enable_x64():
+      self.assertEqual(func().dtype, "float64")
+    with disable_x64():
+      self.assertEqual(func().dtype, "float32")
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_jit={}".format(jit), "jit": jit}
       for jit in ["python", "cpp", None]))
@@ -113,9 +143,11 @@ def func_x64():
       self.assertEqual(x32.result(), jnp.int32)
 
   def test_jit_cache(self):
-    # TODO(jakevdp): enable this test when CPP jit cache is fixed.
-    if FLAGS.experimental_cpp_jit:
-      self.skipTest("Known failure due to https://github.com/google/jax/issues/5532")
+    if jtu.device_under_test() == "tpu":
+      self.skipTest("64-bit random not available on TPU")
+    if jax.lib._xla_extension_version < 4 and FLAGS.experimental_cpp_jit:
+      self.skipTest(
+          "Known failure due to https://github.com/google/jax/issues/5532")
 
     f = partial(random.uniform, random.PRNGKey(0), (1,), 'float64', -1, 1)
     with disable_x64():