pmap complains about shapes even when given bad types

Here’s an example of an error message which seems to obscure the most important and actionable issue:

import jax
import jax.numpy as jnp

class A:
  @jax.pmap
  def f(self, x):
    return x * 2

a = A()
y = a.f(jnp.array([1]))  # prints 'tracing!'

ValueError: pmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())

It’d be better here to mention that A is not a mappable type!

[aaronsarna]

I just want to call out that regardless of what @dwromero 's actual issue is, the fact that the error message provides no useful hints about what the issue is is exactly why this issue was filed and why it should be fixed.

[dwromero]

For those who that might help, my problem was that I forgot to replicate the state across devices. You can do that with state = flax.jax_utils.replicate(state).