Issue when running on multiple devices

[sokol11]

Hi. I tried running some of the examples on a multi-TPU machine and it throws this "assert not ragged" error: jax-ml/jax#13931

The error occurs even when telling jit to use a single device, like so: x, state = jax.jit(strategy.ask, device=jax.devices("tpu")[0])(rng_ask, state)

Running without jit, e.g., x, state = strategy.ask(rng_ask, state) causes input shape errors, which I guess is somewhat expected.

Everything runs fine on a single-device machine.