Repository: google/jax

Issue:
This is a bug in JAX's shape-checking rules; please report it!
I'm getting `This is a bug in JAX's shape-checking rules; please report it!`.

The function that creates this error is the following:

```
@functools.partial(jax.pmap, axis_name="batch")
def accuracy_and_loss_fn2(params, images, labels):
  logits = resnet_apply(params, images)
  loss = cross_entropy_loss(logits=logits, labels=labels)
  accu = jax.lax.pmean(jnp.argmax(logits, axis=1) == labels, axis_name="batch")
  return accu, loss
```

When I replace the `jax.lax.pmean` by a `jnp.mean` it does not cause the error. This is when running in multi-host setup, didn't try single-host yet.

full error:

```
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
```


Discussion:
Here is `cross_entropy_loss`, but `apply_resnet` is too much to post here. It's a ResNet variant which does work in other calls, including multi-host.

```
def cross_entropy_loss(*, logits, labels):
   logp = jax.nn.log_softmax(logits)
   loglik = jnp.take_along_axis(logp, labels[:, None], axis=1)
   return -jnp.mean(loglik)
```

For the offending line, I have also tried `jax.lax.all_gather` as well as `jax.lax.psum` and they result in the same, in case that helps.
Okay cool, presumably if we replace `apply_resnet(...)` with just a constant array of some shape we'll see the same shape error.

It does help that `psum` raises the same error!

One last question, if it's easy to answer: what's the shape of `jnp.argmax(logits, axis=1) == labels` here? We might be able to repro with just that being a constant array, so we can clear out everything except the `psum`/`pmean`.
Sorry, original message was wrong, I was confused. Shape of that is `(B,)` where `B` is the per-device batch-size, for example `(2,)` in a recent run I did.

~It *should* be `(B, C)` where `B` is batch-size and `C = 1000` in my case, but I will add a logging and run again just to be sure!~
Argh, miss-clicked.
Perfect, ty!
This error looks like it's saying we're creating an XLA add instruction on boolean inputs (which isn't supported). (But it's evidently not going through the same code path as `jax.lax.add(jax.numpy.array([False]), jax.numpy.array([True]))`, which throws a better error.)
@skye had a really smart guess: we lower `jax.core.unit` to `Pred[]`, so maybe there's a partial-eval-inserted unit cropping up here. (It'd be nice to revise things so we lower units to nothing at all in HLO... one of the many things on the todo list.)

I haven't checked yet whether this is pod-specific; on one hand it doesn't sound pod-specific, on the other hand I'm sure we have test coverage for this kind of use case...
@jekbradbury is right, in this case a boolean array is being passed to pmean. I think we should improve the error message, but still make it an error and require an explicit cast in this situation. However, there is an argument to be made for providing some kind of implicit casting functionality (either in psum or even a separate API endpoint?), since that's what numpy does, and even though psum is in lax there isn't a numpy equivalent.

Any thoughts? Otherwise I'll just improve the error message for now.
Oh, hah, I totally didn't get what @jekbradbury 's point was. Sorry! I get it now.

I think we should handle booleans (by promoting to a non-boolean type, like we do for analogous situations in lax_numpy.py).
Ok! I actually like that better, it just felt weird to have some lax primitives do casting and others not. But this way is more useful!
Yeah, it is a bit weird, but I think most lax primitives aren't user-facing. The bool->int promotion in a reduce-sum seems pretty harmless, and this is a nice canonical use case too.
(Also, as per the title of this issue, actually we should do dtype checking in the abstract eval / shape rule, but if we do any it's not currently matching XLA's constraints.)
Wait are you sure the dtype isn't a red herring here? As I mentioned, I get the same error if I do `jax.lax.all_gather` instead of `psum`, which I should reasonably expect to work, no?

Regarding casting, my understanding is "jax is numpy" and so I would expect implicit cast of bool to float, which I use all the time in numpy.
`jax.numpy` approximates NumPy semantics. `jax.lax` in general does not. Although there's certainly a case to be made that the user-facing parallel operations should promote like NumPy, even if other `lax` operations do not.
@lucasb-eyer I'm surprised you're getting exactly the same error with an all-gather, since the error specifically talks about add: `Expected element type in shape to be arithmetic type for operation add; got PRED.` I haven't had a chance to dig into this myself with your code yet though.

Either way, I was able to repro this error with a small psum(bool) unit test, so there's definitely an issue there. I was gonna start by fixing the unit test, then rerunning your code with the fix and see what happens. It's also possible the unit test will expose the more general issue if there is one.
[`all_gather` is just a psum](https://github.com/google/jax/blob/77e31323f79d01b9e97b52546571b6c98dc4df2f/jax/lax/lax_parallel.py#L460) so it makes sense it'd be the same error.

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -44,6 +44,8 @@ def psum(x, axis_name, *, axis_index_groups=None):
   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
@@ -68,10 +70,13 @@ def psum(x, axis_name, *, axis_index_groups=None):
   >>> 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))
+  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)
 
 def pmean(x, axis_name, *, axis_index_groups=None):
   """Compute an all-reduce mean on ``x`` over the pmapped axis ``axis_name``.


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1184,6 +1184,26 @@ def matrix_vector(x, y, parallel=True):
     self.assertAllClose(result1, result3, check_dtypes=False, atol=1e-3, rtol=1e-3)
     self.assertAllClose(result1, result4, check_dtypes=False, atol=1e-3, rtol=1e-3)
 
+  def testPsumOnBooleanDtype(self):
+    # https://github.com/google/jax/issues/3123
+    n = xla_bridge.device_count()
+    if n > 1:
+      x = jnp.array([True, False])
+
+      out = pmap(lambda x: jax.lax.psum(x, 'i'), 'i')(x)
+      self.assertEqual(list(out), [1, 1])
+
+      out = pmap(lambda x: jax.lax.pmean(x, 'i'), 'i')(x)
+      self.assertEqual(list(out), [1/2, 1/2])
+    else:
+      x = jnp.array([True])
+
+      out = pmap(lambda x: jax.lax.psum(x, 'i'), 'i')(x)
+      self.assertEqual(list(out), [1])
+
+      out = pmap(lambda x: jax.lax.pmean(x, 'i'), 'i')(x)
+      self.assertEqual(list(out), [1])
+
 
 class PmapWithDevicesTest(jtu.JaxTestCase):