Encountered AllReduce on 3 replicas, which is odd and greater than one. This case is not implemented for TPU.

[vwxyzjn]

Hi JAX team! Can I do jax.lax.psum on 3 TPU devices? The following code passes for GPU devices. I didn't find any issues by googling, so just thought I might post a question here :)

import jax
import jax.numpy as jnp
from jax import pmap

# Define a function to be parallelized
def my_function(x):
    return jax.lax.psum(x, 'i')

# Define an input array
x = jnp.arange(3)

print(jax.devices())
parallel_fn = pmap(my_function, devices=jax.devices()[:3], axis_name='i')
result = parallel_fn(x)
print(result)

    compiled = dispatch.compile_or_get_cached(
  File "/home/costahuang/.cache/pypoetry/virtualenvs/cleanba-pOszUjJ6-py3.9/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1077, in compile_or_get_cached
    return backend_compile(backend, serialized_computation, compile_options,
  File "/home/costahuang/.cache/pypoetry/virtualenvs/cleanba-pOszUjJ6-py3.9/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/costahuang/.cache/pypoetry/virtualenvs/cleanba-pOszUjJ6-py3.9/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1012, in backend_compile
    return backend.compile(built_c, compile_options=options)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: Encountered AllReduce on 3 replicas, which is odd and greater than one. This case is not implemented for TPU.

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 "/home/costahuang/cleanba/test.py", line 14, in <module>
    result = parallel_fn(x)
jaxlib.xla_extension.XlaRuntimeError: UNIMPLEMENTED: Encountered AllReduce on 3 replicas, which is odd and greater than one. This case is not implemented for TPU.