Repository: google/jax

Issue:
Breaking changes on vjp with with tokens for `jax>=0.2.8`?
After updating jax to the latest version I'm starting to see test failures on (mpi4jax)[https://github.com/PhilipVinc/mpi4jax] when calling `vjp` on functions that accept tokens.

see below:

```python
def test_allreduce_vjp():
    from mpi4jax import Allreduce

    token = jax.lax.create_token(123)
    arr = jnp.ones((3, 2))
    _arr = arr.copy()

    res, vjp_fun = jax.vjp(lambda x: Allreduce(x, op=MPI.SUM, token=token)[0], arr)
    (vjp,) = vjp_fun(_arr)

    expected, _ = Allreduce(arr, op=MPI.SUM)
    assert jnp.array_equal(expected, res)
    assert jnp.array_equal(_arr, vjp)

```

fails with error
```python
tests/test_collective_ops.py:202:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/api.py:1866: in vjp
    return _vjp(lu.wrap_init(fun), *primals, has_aux=has_aux)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/api.py:1874: in _vjp
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/interpreters/ad.py:114: in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/interpreters/ad.py:101: in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:506: in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/linear_util.py:160: in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
tests/test_collective_ops.py:202: in <lambda>
    res, vjp_fun = jax.vjp(lambda x: Allreduce(x, op=MPI.SUM, token=token)[0], arr)
mpi4jax/validation.py:90: in wrapped
    return function(*args, **kwargs)
mpi4jax/collective_ops/allreduce.py:68: in Allreduce
    return mpi_allreduce_p.bind(x, token, op=op, comm=comm, transpose=_transpose)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/core.py:270: in bind
    tracers = map(top_trace.full_raise, args)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/_src/util.py:37: in safe_map
    return list(map(f, *args))
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/core.py:377: in full_raise
    return self.pure(val)
../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/interpreters/ad.py:269: in pure
    tangent_zero = Zero(get_aval(val).at_least_vspace())
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = AbstractToken()

    def at_least_vspace(self):
>     raise NotImplementedError("must override")
E     NotImplementedError: must override

../../../../../Documents/pythonenvs/mpi4jax_env/lib/python3.8/site-packages/jax/core.py:803: NotImplementedError
```

In mpi4jax, the culprit seems to be 
`return mpi_allreduce_p.bind(x, token, op=op, comm=comm, transpose=_transpose)`, which normally works fine, also under jit, but if under a vjp transformation he does not like the token to be there?

This code was working with a previous version of jax.


Discussion:
This test was successful with `jax==0.2.7` and is broken with `jax>=0.2.8`

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -1134,6 +1134,7 @@ def join(self, other):
     else:
       assert False, f"Cannot join {self} with {other}"
   def str_short(self): return 'Tok'
+  def at_least_vspace(self): return self
 
 abstract_token = AbstractToken()
 
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -367,10 +367,9 @@ def _execute_replicated_primitive(prim, compiled, result_handler, *args):
 
 
 def check_special(prim, bufs):
-  for buf in bufs:
-    # TODO(jblespiau): We can simply use buf.xla_shape() when version 0.1.58 is
-    # the default.
-    _check_special(prim.name, getattr(buf, "xla_shape", buf.shape)(), buf)
+  if FLAGS.jax_debug_infs or FLAGS.jax_debug_nans:
+    for buf in bufs:
+      _check_special(prim.name, buf.xla_shape(), buf)
 
 def _check_special(name, xla_shape, buf):
   assert not xla_shape.is_tuple()


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2160,6 +2160,20 @@ def test_linearize_aval_error(self):
     with self.assertRaisesRegex(ValueError, "tangent values inconsistent"):
       f_jvp(np.ones(2, np.int32))
 
+  def test_grad_of_token_consuming_primitive(self):
+    # https://github.com/google/jax/issues/5463
+    tokentest_p = core.Primitive("tokentest")
+    tokentest_p.def_impl(partial(xla.apply_primitive, tokentest_p))
+    tokentest_p.def_abstract_eval(lambda x, y: x)
+    xla.translations[tokentest_p] = lambda c, x, y:  x
+    ad.defjvp(tokentest_p, (lambda g, x, token: x), None)
+
+    token = jax.lax.create_token(123)
+    arr = jnp.ones((3, 2))
+    res, vjp_fun = jax.vjp(lambda x: tokentest_p.bind(x, token), arr)
+    # Should not crash.
+    vjp_fun(arr)
+
 
 class RematTest(jtu.JaxTestCase):