Repository: google/jax

Issue:
The side-effect omnistaging example is broken
The example illustrating a possible issue with omnistaging and global state (https://github.com/google/jax/blob/master/design_notes/omnistaging.md#example-1) no longer leads to an escaped tracer:
```
from jax import jit
from jax import random

print('omnistaging enabled', jax.config.omnistaging_enabled)
print(jax.__version__)

key = random.PRNGKey(0)
def init():
  global key
  key, subkey = random.split(key)
  print(key)
  return random.normal(subkey, ())
print(init())  # -1.2515389
print(init())  # -0.58665067
init = jit(init)
print(init())  # 0.48648298
print(init())  # 0.48648298  !!
print(key) # Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
print(random.normal(key, ()))

# omnistaging enabled True
# 0.2.7
# [4146024105  967050713]
# -1.2515285
# [2384771982 3928867769]
# -0.5866531
# [3382499631 3878610767]
# 0.48647928
# 0.48647928
# [3382499631 3878610767]
# Buffer(-0.4876217, dtype=float32)
```
Has this omnistaging behaviour been changed?


Discussion:
Thanks for flagging this - I tried back as far as jax 0.2.0 when omnistaging was first turned on by default, and could not reproduce the expected escaped tracer error. I'm going to assign to @mattjj because he might know what's going on.
Hrm, curious. This repros for me:

```python
import jax
from jax import jit
from jax import random

key = random.PRNGKey(0)
def init():
  global key
  key, subkey = random.split(key)
  return random.normal(subkey, ())
init = jit(init)
print(init())  # -1.2515389
print(init())  # -1.2515389
print(key) # Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
```

If we put a single `print(init())` before jitting, then we also get the escaped tracer. But if we put two, we don't! Very curious.

I think the reason why @jakevdp wasn't able to reproduce is that he didn't rewind jaxlib as well as jax. The funny actor here is the C++ jit. If we run the original version with `JAX_CPP_JIT=0` then we get the behavior as expected:

```
$ cat 5206.py
import jax
from jax import jit
from jax import random

key = random.PRNGKey(0)
def init():
  global key
  key, subkey = random.split(key)
  return random.normal(subkey, ())
print(init())
print(init())
init = jit(init)
print(init())
print(init())
print(key)

$ python 5206.py
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
-1.2515389
-0.58665067
0.48648298
0.48648298
[3382499631 3878610767]

$ env JAX_CPP_JIT=0 python 5206.py
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
-1.2515389
-0.58665067
0.48648298
0.48648298
Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
```

I was tipped off because in the last line of the OP we see a `Buffer` printed out, and that smelled like a C++ issue.

I don't understand this in full yet though. cc @jblespiau in case he has ideas.
Just to underscore, I really don't know what's going on here, and I don't mean to imply that any particular code has a bug. Setting `JAX_CPP_JIT` affects the behavior, but I don't actually know if there's an issue in the C++ code, or related Python code, or if actually this is somehow intended behavior as a consequence of some of our caching decisions related to the C++ jit.
```
key = random.PRNGKey(0)
def init(a):
  global key
  import pdb; pdb.set_trace()
  key, subkey = random.split(key)
  return random.normal(subkey, ())
print(init(1))
print(init(2))
init = jit(init)
print(init(3))
print(init(4))
print(key)
```

What happens?

print(init()) << This compiles `random.split`. at this moment, random.split is provided with a tracer, it compiles, run the execution and return concrete keys.
print(key) --> prints  [4146024105  967050713]
print(init()) << Here, we hit the cache for these, it returns a real value, as expected
print(key) --> [2384771982 3928867769]
init = jit(init)
print(init()) <-- This triggers the compilation of `init`. `ramdom.split` is provided with a concrete key so, it will return a concrete type
print(key) >> [3382499631 3878610767]
print(init()) <-- No longer executing Python anyway. key won't change.
print(key) >> [3382499631 3878610767]


It you want to get an abstract type, you need to store it when there is one (e.g. if you add an argument "a", and you save it, it will work:

```
global_v = None
def f(a):
  global global_v
  global_v = a
  return a + 1

jit(f)(1)
print(global_v)
```

So for the C++, I can explain.

However, for the Python I cannot.

When we start tracing, do we set a global variable somewhere, to tell to compiled function they should trace?

i.e. when we jit `init`, and we get into `random.split`, it it tracing, even though we passed concrete types.

If we think functional (i.e. same inputs, same outputs), this seems to be the correct behavior (random.split is given concrete types, returns concrete types), while the Python execution is not functional, as its behavior changes based on some external values).

Ah yes, this is the nature of omnistaging, namely that with omnistaging things like `jit` change the behavior of operations based on global context, not just on their input arguments. That's the goal: not to rely solely on data dependence, so that we can stage out more, like constant creation and RNG operations. The cost is that JAX's tracing behavior itself isn't purely functional, in that that whether a function produces a tracer or not cannot be decided based on whether the inputs are traced.

We shouldn't get a cache hit on the `print(init())` line after applying `jit`. The reason is that it results in a worse program being staged out, at least if we imagine the random constant being arbitrarily large: there's a potentially-large random value staged out into the jitted version of `init`, as opposed to the intended behavior of staging out the PRNG generation itself (based on the small key as a constant).

But I see now why we are getting that behavior: it's because, by design, with the C++ jit we've moved the cache to the api.py level, whereas with the Python version we have it in the jit impl (i.e. the xla_call impl). In the Python version, when running the jitted version of `init`, we never hit the impl rule for `random.split` at all due to the change in global context, so there's no risk of hitting the cache. That is, the Python version is already sensitive to the necessary global context.

I think we need to make the C++ jit sensitive to the global context. Otherwise, we're not getting the full benefits of omnistaging. Concretely, I think we need to make the C++ cache depend on the value of the hashable `core.thread_local_state.trace_state.trace_stack.dynamic`.
(Thanks for catching this and letting us know, @sbodenstein ! Otherwise the ill effects are subtle, like increased memory fragmentation, and we might have been scratching our heads over the reason.)
> مثالی که یک مسئله احتمالی با وضعیت جهانی و جهانی را نشان می دهد ( https://github.com/google/jax/blob/master/design_notes/omnistaging.md#example-1 ) دیگر منجر به ردیابی فرار نمی شود:
> 
> ```
> from jax import jit
> from jax import random
> 
> print('omnistaging enabled', jax.config.omnistaging_enabled)
> print(jax.__version__)
> 
> key = random.PRNGKey(0)
> def init():
>   global key
>   key, subkey = random.split(key)
>   print(key)
>   return random.normal(subkey, ())
> print(init())  # -1.2515389
> print(init())  # -0.58665067
> init = jit(init)
> print(init())  # 0.48648298
> print(init())  # 0.48648298  !!
> print(key) # Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
> print(random.normal(key, ()))
> 
> # omnistaging enabled True
> # 0.2.7
> # [4146024105  967050713]
> # -1.2515285
> # [2384771982 3928867769]
> # -0.5866531
> # [3382499631 3878610767]
> # 0.48647928
> # 0.48647928
> # [3382499631 3878610767]
> # Buffer(-0.4876217, dtype=float32)
> ```
> 
> آیا این رفتار همه کاره تغییر کرده است؟

Ok

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -259,7 +259,7 @@ def _cpp_jit(
     raise ValueError("can't specify both a device and a backend for jit, "
                      f"got device={device} and backend={backend}.")
 
-  def cache_miss(*args, **kwargs):
+  def cache_miss(_, *args, **kwargs):
     ### This first part is basically the same code as in _python_jit.
     # An alternative would be for cache_miss to accept from C++ the arguments
     # (dyn_args, donated_invars, args_flat, in_tree), since otherwise we have
@@ -362,17 +362,20 @@ def get_jax_disable_jit_flag():
     """
     return config.read("jax_disable_jit")
 
+  static_argnums_ = (0,) + tuple(i + 1 for i in static_argnums)
   cpp_jitted_f = jax_jit.jit(fun, cache_miss, get_device_info,
                              get_jax_enable_x64, get_jax_disable_jit_flag,
-                             static_argnums)
+                             static_argnums_)
 
   # TODO(mattjj): make cpp callable follow descriptor protocol for bound methods
   @wraps(fun)
   @api_boundary
   def f_jitted(*args, **kwargs):
+    context = getattr(core.thread_local_state.trace_state.trace_stack,
+                      'dynamic', None)
     # TODO(jblespiau): Move this to C++.
     if FLAGS.jax_debug_nans and not _jit_is_disabled():
-      device_arrays = cpp_jitted_f(*args, **kwargs)
+      device_arrays = cpp_jitted_f(context, *args, **kwargs)
       try:
         xla.check_nans(xla.xla_call_p, [
             da.device_buffer
@@ -384,9 +387,11 @@ def f_jitted(*args, **kwargs):
         assert FLAGS.jax_debug_nans  # compiled_fun can only raise in this case
         print("Invalid nan value encountered in the output of a C++-jit "
               "function. Calling the de-optimized version.")
-        return cache_miss(*args, **kwargs)[0]  # probably won't return
-    else:
+        return cache_miss(context, *args, **kwargs)[0]  # probably won't return
+    elif _jit_is_disabled():
       return cpp_jitted_f(*args, **kwargs)
+    else:
+      return cpp_jitted_f(context, *args, **kwargs)
   f_jitted._cpp_jitted_f = cpp_jitted_f
 
   return f_jitted


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -506,6 +506,22 @@ def f(x):
         f"explicit inner-jit backend specification cpu."):
       f(1.)
 
+  def test_omnistaging(self):
+    # See https://github.com/google/jax/issues/5206
+    if not config.omnistaging_enabled:
+      raise unittest.SkipTest("test only works with omnistaging")
+
+    key_list = [None]
+
+    def init():
+      key, subkey = jax.random.split(key_list[0])
+      key_list[0] = key
+      return jax.random.normal(subkey, ())
+
+    key_list[0] = np.array([2384771982, 3928867769], dtype=np.uint32)
+    init()
+    self.jit(init)()
+    self.assertIsInstance(key_list[0], core.Tracer)
 
 class PythonJitTest(CPPJitTest):