Understanding a JAX error message from jited function

[luyug]

I've got the following message when running a jitted function on 2 GPUs.

jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Executable expected shape f32[2048,512]{1,0} for argument 3 but got incompatible shape f32[1024,512]{1,0}: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

I am wondering what this means/entails?

[jakevdp]

Thanks for the report. Can you provide a code sample that reproduces this issue? This is an internal error from XLA (and definitely represents a JAX bug) but it's hard to say much else from the error alone.

[luyug]

Thanks for the quick reply! My current code base is kind of convoluted. Let me try to extract a minimal piece for repro.

My quick look seems to suggest that it could be some weird interactions between the GPU partitioner and jax's einsum. Anyway, let me work out the repro.