# Buggy code
def psum(x, axis_name, *, axis_index_groups=None):
    """Compute an all-reduce sum on ``x`` over the pmapped axis ``axis_name``.

    If ``x`` is a pytree then the result is equivalent to mapping this function to
    each leaf in the tree.

    Args:
      x: array(s) with a mapped axis named ``axis_name``.
      axis_name: hashable Python object used to name a pmapped axis (see the
        ``pmap`` docstring for more details).
      axis_index_groups: optional list of lists containing axis indices (e.g. for
        an axis of size 4, [[0, 1], [2, 3]] would perform psums over the first
        two and last two replicas). Groups must cover all axis indices exactly
        once, and all groups must be the same size.


    Returns:
      Array(s) with the same shape as ``x`` representing the result of an
      all-reduce sum along the axis ``axis_name``.

    For example, with 4 XLA devices available:

    >>> x = np.arange(4)
    >>> y = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(x)
    >>> print(y)
    [6 6 6 6]
    >>> y = jax.pmap(lambda x: x / jax.lax.psum(x, 'i'), axis_name='i')(x)
    >>> print(y)
    [ 0.          0.16666667  0.33333334  0.5       ]
    """
    leaves, treedef = tree_util.tree_flatten(x)
    _validate_axis_index_groups(axis_index_groups)
    return treedef.unflatten(
        psum_p.bind(*leaves, axis_name=axis_name, axis_index_groups=axis_index_groups)
    )


# Traceback (RuntimeError)
Traceback (most recent call last):
File "[REDACTED]/py/jax/interpreters/xla.py", line 301, in primitive_computation
return c.build()
RuntimeError: Invalid argument: Expected element type in shape to be arithmetic type for operation add; got PRED.:     @     0x55b73da20345  xla::XlaBuilder::BinaryOp()
@     0x55b73da19dc9  xla::Add()
@     0x7fc98c23d5f4  pybind11::cpp_function::initialize<>()::{lambda()#1}::__invoke()
@     0x7fc98c1e9963  pybind11::cpp_function::dispatcher()
@     0x55b745431fe7  PyCFunction_Call
@     0x55b7454b3adf  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450f983e  partial_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7454b3a04  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450fa680  bounded_lru_cache_wrapper
@     0x55b7453e64ba  PyObject_Call
@     0x55b7454b3a04  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454b7a24  fast_function
@     0x55b7454b66bc  call_function
@     0x55b7454b375f  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450f983e  partial_call
@     0x55b7453e678f  _PyObject_FastCallDict
@     0x55b7454b669d  call_function
@     0x55b7454b375f  _PyEval_EvalFrameDefault
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "[REDACTED]/py/absl/app.py", line 464, in run
_run_main(main, args)
File "[REDACTED]/py/absl/app.py", line 393, in _run_main
sys.exit(main(argv))
File "[REDACTED]/imagenet/train.py", line 255, in main
batch["image"], batch["label"])
File "[REDACTED]/py/jax/api.py", line 1051, in f_pmapped
mapped_invars=tuple(axis is not None for axis in in_axes_flat))
File "[REDACTED]/py/jax/core.py", line 1021, in _call_bind
outs = primitive.impl(f, *args, **params)
File "[REDACTED]/py/jax/interpreters/pxla.py", line 604, in xla_pmap_impl
*abstract_args)
File "[REDACTED]/py/jax/linear_util.py", line 221, in memoized_fun
ans = call(fun, *args)
File "[REDACTED]/py/jax/interpreters/pxla.py", line 702, in parallel_callable
extend_name_stack(wrap_name(name, 'pmap')), *xla_args)
File "[REDACTED]/py/jax/interpreters/xla.py", line 406, in jaxpr_subcomp
**new_params)
File "[REDACTED]/py/jax/lax/lax_parallel.py", line 311, in _psum_translation_rule
return _notuple_psum_translation_rule(c, *args, replica_groups=replica_groups)
File "[REDACTED]/py/jax/lax/lax_parallel.py", line 357, in _notuple_psum_translation_rule
return xops.Tuple(c, list(map(_translate, args)))
File "[REDACTED]/py/jax/lax/lax_parallel.py", line 356, in _translate
return psum(val)
File "[REDACTED]/py/jax/lax/lax_parallel.py", line 304, in _allreduce_translation_rule
computation = xla.primitive_subcomputation(prim, scalar, scalar)
File "[REDACTED]/py/jax/interpreters/xla.py", line 309, in primitive_subcomputation
return primitive_computation(prim, AxisEnv(1), None, False, *avals, **params)
File "[REDACTED]/py/jax/interpreters/xla.py", line 306, in primitive_computation
raise RuntimeError(msg) from e
RuntimeError: Invalid argument: Expected element type in shape to be arithmetic type for operation add; got PRED.:     @     0x55b73da20345  xla::XlaBuilder::BinaryOp()
@     0x55b73da19dc9  xla::Add()
@     0x7fc98c23d5f4  pybind11::cpp_function::initialize<>()::{lambda()#1}::__invoke()
@     0x7fc98c1e9963  pybind11::cpp_function::dispatcher()
@     0x55b745431fe7  PyCFunction_Call
@     0x55b7454b3adf  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450f983e  partial_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7454b3a04  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450fa680  bounded_lru_cache_wrapper
@     0x55b7453e64ba  PyObject_Call
@     0x55b7454b3a04  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454b7a24  fast_function
@     0x55b7454b66bc  call_function
@     0x55b7454b375f  _PyEval_EvalFrameDefault
@     0x55b7454b7160  _PyEval_EvalCodeWithName
@     0x55b7454ad938  PyEval_EvalCodeEx
@     0x55b7454171a6  function_call
@     0x55b7453e64ba  PyObject_Call
@     0x55b7450f983e  partial_call
@     0x55b7453e678f  _PyObject_FastCallDict
@     0x55b7454b669d  call_function
@     0x55b7454b375f  _PyEval_EvalFrameDefault
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues

# Fixed code
def psum(x, axis_name, *, axis_index_groups=None):
    """Compute an all-reduce sum on ``x`` over the pmapped axis ``axis_name``.

    If ``x`` is a pytree then the result is equivalent to mapping this function to
    each leaf in the tree.

    Inputs of boolean dtype are converted to integers before the reduction.

    Args:
      x: array(s) with a mapped axis named ``axis_name``.
      axis_name: hashable Python object used to name a pmapped axis (see the
        ``pmap`` docstring for more details).
      axis_index_groups: optional list of lists containing axis indices (e.g. for
        an axis of size 4, [[0, 1], [2, 3]] would perform psums over the first
        two and last two replicas). Groups must cover all axis indices exactly
        once, and all groups must be the same size.


    Returns:
      Array(s) with the same shape as ``x`` representing the result of an
      all-reduce sum along the axis ``axis_name``.

    For example, with 4 XLA devices available:

    >>> x = np.arange(4)
    >>> y = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(x)
    >>> print(y)
    [6 6 6 6]
    >>> y = jax.pmap(lambda x: x / jax.lax.psum(x, 'i'), axis_name='i')(x)
    >>> print(y)
    [ 0.          0.16666667  0.33333334  0.5       ]
    """
    _validate_axis_index_groups(axis_index_groups)
    leaves, treedef = tree_util.tree_flatten(x)
    leaves = [
        lax.convert_element_type(l, onp.int32) if dtypes.dtype(l) == onp.bool_ else l
        for l in leaves
    ]
    out_flat = psum_p.bind(
        *leaves, axis_name=axis_name, axis_index_groups=axis_index_groups
    )
    return tree_util.tree_unflatten(treedef, out_flat)