Repository: google/jax

Issue:
JITing no-ops breaks with tokens
Reproducer:

```python
>>> import jax

>>> @jax.jit
>>> def noop(arr, token):
...    return arr, token

>>> arr = jax.numpy.ones(10)
>>> token = jax.lax.create_token()

>>> noop(arr, token)
Traceback (most recent call last):
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 119, in device_put
    return device_put_handlers[type(x)](x, device)
KeyError: <class 'jax.interpreters.xla.Token'>

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 143, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/_src/api.py", line 426, in cache_miss
    out_flat = xla.xla_call(
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/core.py", line 1565, in bind
    return call_bind(self, fun, *args, **params)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/core.py", line 1556, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/core.py", line 1568, in process
    return trace.process_call(self, fun, tracers, params)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/core.py", line 609, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 581, in _xla_call_impl
    return compiled_fun(*args)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 905, in _execute_trivial
    return [_copy_device_array_to_device(x, device) if type_is_device_array(x)
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 906, in <listcomp>
    else h(*device_put(x, device)) for h, x in zip(handlers, outs)]
  File "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 121, in device_put
    raise TypeError(f"No device_put handler for type: {type(x)}") from err
jax._src.traceback_util.UnfilteredStackTrace: TypeError: No device_put handler for type: <class 'jax.interpreters.xla.Token'>

The stack trace below excludes JAX-internal frames.
The preceding 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 "/groups/ocean/dhaefner/miniconda3/envs/veros/lib/python3.8/site-packages/jax/interpreters/xla.py", line 121, in device_put
    raise TypeError(f"No device_put handler for type: {type(x)}") from err
TypeError: No device_put handler for type: <class 'jax.interpreters.xla.Token'>
```


Discussion:

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -1043,6 +1043,7 @@ class Token(object): pass
 xla_shape_handlers[AbstractToken] = lambda _: (xc.Shape.token_shape(),)
 xla_result_handlers[AbstractToken] = lambda _, __: lambda _: token
 canonicalize_dtype_handlers[Token] = identity
+device_put_handlers[Token] = lambda x, _: (x,)
 
 
 def _forward_method(attrname, self, fun, *args):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -369,6 +369,16 @@ def test_trivial_computations(self):
     self.assertIs(z3, x1)
     self.assertEqual(z2, 1)
 
+  def test_trivial_computations_with_tokens(self):
+    @self.jit
+    def noop(arr, token):
+      return arr, token
+
+    arr = jax.numpy.ones(10)
+    token = jax.lax.create_token()
+
+    self.assertEqual(token, noop(arr, token)[1])
+
   def test_jit_bad_input(self):
     def f(x):
       return x