Repository: google/jax

Issue:
Cannot jit GPU functions returning Token objects
Cannot jit GPU functions that return Token objects, as reported in #5707 

Seen with  `jax==0.2.9` and `libjax==0.1.61+CUDA111`.
The same snippet works fine when running on CPU (`jax.config.update('jax_platform_name', 'cpu')`).

Reproducer:
```python
>>> import jax
>>> jax.devices()
[GpuDevice(id=0)]
>>> jax.jit(jax.lax.create_token)(1.0)
```

It leads to the error:
```python
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Unimplemented: Unimplemented primitive type TOKEN

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 "<stdin>", line 1, in <module>
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/api.py", line 398, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/api.py", line 289, in cache_miss
    out_flat = xla.xla_call(
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/core.py", line 1275, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/core.py", line 1266, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/core.py", line 1278, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/core.py", line 631, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/interpreters/xla.py", line 580, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/linear_util.py", line 260, in memoized_fun
    ans = call(fun, *args)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/interpreters/xla.py", line 714, in _xla_callable
    donated_invars = set_up_aliases(c, xla_args, out_tuple, donated_invars, tuple_args)
  File "/home/filippovicentini/Documents/pythonenvs/mpi4jax_env/lib64/python3.8/site-packages/jax/interpreters/xla.py", line 753, in set_up_aliases
    key = (element.dimensions(), element.numpy_dtype())
RuntimeError: Unimplemented: Unimplemented primitive type TOKEN
```



Discussion:
Seems to be a followup from #5707

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -744,7 +744,7 @@ def set_up_aliases(c, xla_args, out_tuple, donated_args, tuple_args):
   for arg_index, arg in enumerate(xla_args):
     if donated_args[arg_index]:
       for param_index, element in flatten_shape(c.GetShape(arg)):
-        key = (element.dimensions(), element.numpy_dtype())
+        key = (element.dimensions(), element.xla_element_type())
         if tuple_args:
           param_number = 0
           param_index = (arg_index,) + tuple(param_index)
@@ -756,7 +756,7 @@ def set_up_aliases(c, xla_args, out_tuple, donated_args, tuple_args):
   # Consume donations for outputs.
   out_donated_args = list(donated_args)
   for output_index, element in flatten_shape(c.GetShape(out_tuple)):
-    key = (element.dimensions(), element.numpy_dtype())
+    key = (element.dimensions(), element.xla_element_type())
     if donations.get(key, ()):
       param_number, param_index, arg_index = donations[key].popleft()
       out_donated_args[arg_index] = False


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2216,6 +2216,10 @@ def test_grad_of_token_consuming_primitive(self):
     # Should not crash.
     vjp_fun(arr)
 
+  def test_jit_returning_token(self):
+    x = jax.jit(jax.lax.create_token)(1.0)
+    self.assertIsInstance(x, jax.interpreters.xla.Token)
+
   def test_leak_checker_catches_a_jit_leak(self):
     if not config.omnistaging_enabled:
       raise unittest.SkipTest("test only works with omnistaging")