flax_on_pjit.ipynb's pjit train_step reshards sharded data? Tut. crashes in multi-host SPMD environment.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
RHEL 8.5, Docker image: cuda_11.5.2-cudnn8-devel-ubuntu20.04 (tried cuda_11.8.0-cudnn8-devel-ubuntu20.04)

- Flax, jax, jaxlib versions:
flax==0.6.5 (tried 0.6.4)
jax==0.4.4
jaxlibb==0.4.4+cuda11.cudnn82

- Python version:
3.9.9

- GPU/TPU model and memory:
20x A100 80GB (1 GPU per host)

- CUDA version (if applicable):
cuda/11.5.1-75

### Problem you have encountered:
Using pjit_step_fn repeatedly (i.e.: training) from [flax_on_pjit.ipynb](https://colab.research.google.com/github/google/flax/blob/master/docs/guides/flax_on_pjit.ipynb) causes a crash when used with global devices.

The tutorial [creates a pjit_step_fn here](https://colab.research.google.com/github/google/flax/blob/master/docs/guides/flax_on_pjit.ipynb#scrollTo=4e3cc300cfee&line=16&uniqifier=1).

I've tried this with a dict containing the data and passing the jax.Array in as the original code does - same issue.
I have also tried a simpler model with just one axis constraint.

### What you expected to happen:
The tutorial says it is also for multi-host environments, so I would have expected the provided train function to work here with a multi-host environment.
If this method is not suitable, can you please point me in the direction of what I should be using? mesh_utils' global broadcasts are much too slow for getting the mean of the gradients.
For what it's worth, PMAP is working globally, as expected.

### Logs, error messages, etc:

With 20 devices, a mesh shape of (10, 2), data shaped (20, 1000), pjit_step_fn seems to shard what is already sharded?

```
2023-02-23 15:35:49.935174: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2410] Execution of replica 0 failed: INVALID_ARGUMENT: Executable expected shape f32[2,1000]{1,0} for argument 8 but got incompatible shape f32[20,1000]{1,0}
ValueError: INVALID_ARGUMENT: Executable expected shape f32[2,1000]{1,0} for argument 8 but got incompatible shape f32[20,1000]{1,0}
```

### Steps to reproduce:
It's just [flax_on_pjit.ipynb](https://colab.research.google.com/github/google/flax/blob/master/docs/guides/flax_on_pjit.ipynb), but with jax.distributed.initialize() added to get global devices.
This code causes the crash, by trying to use the updated state from pjit_step_fn. The first call will not crash.
https://colab.research.google.com/drive/1F2aOQWXo0O1e8dezXjjyN302TZBkV_Lx?usp=sharing