create_token XLA error on GPU (and rarely CPU)

I ran example from readme on version 0.2.7 (installation for other versions fails because of this issue). Below is my code and error: Am I doing something wrong?

code:

import jax
import mpi4jax

comm = MPI.COMM_WORLD
a = jax.numpy.ones((5,4))
b = mpi4jax.Allreduce(a, op=MPI.SUM, comm=comm)
b_jit = jax.jit(lambda x: mpi4jax.Allreduce(x, op=MPI.SUM, comm=comm))(a)

error:

  File "test.py", line 7, in <module>
    b = mpi4jax.Allreduce(a, op=MPI.SUM, comm=comm)
  File "/opt/conda/lib/python3.7/site-packages/mpi4jax/validation.py", line 90, in wrapped
    return function(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/mpi4jax/collective_ops/allreduce.py", line 66, in Allreduce
    return mpi_allreduce_p.bind(x, token, op=op, comm=comm, transpose=_transpose)
  File "/opt/conda/lib/python3.7/site-packages/jax/core.py", line 282, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/opt/conda/lib/python3.7/site-packages/jax/core.py", line 628, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py", line 238, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
  File "/opt/conda/lib/python3.7/site-packages/jax/_src/util.py", line 198, in wrapper
    return cached(bool(FLAGS.jax_enable_x64), *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/jax/_src/util.py", line 191, in cached
    return f(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py", line 288, in xla_primitive_callable
    compiled = backend_compile(backend, built_c, options)
  File "/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py", line 352, in backend_compile
    return backend.compile(built_c, compile_options=options)
RuntimeError: Internal: Unexpected shape kind for %after-all.2 = token[] after-all(), metadata={op_type="allreduce_mpi" op_name="allreduce_mpi[ comm=<mpi4jax.utils.HashableMPIType object at 0x7fd1025dfa20>\n               op=<mpi4jax.utils.HashableMPIType object at 0x7fd1025df9b0>\n               transpose=False ]"} and shape index {}

[PhilipVinc]

More probably it’s jax starting to validate token data, which we do not properly propagate along our operations.

I opened Jax#5707 for that.

@fgvbrt I’ll try to quickly fix mpi4jax once jax people tell us how to treat this case. in the meantime, the only workaround is to use an older version of jax, libjax and mpi4jax.

[dionhaefner]

Fixed in the new jaxlib release.