Confusion around Multi-host data loading with jit

[fattorib]

Hi,

I'm currently porting over some old pre-jax.Array pjit data and tensor-parallel code to work with newer JAX versions. I read this section from the Array migration guide and for data batches I can wrap the batch in host_local_array_to_global_array to recover correct global mesh sharding, as previously pjit only required the batch to be sharded over the submesh local to the host.

However, if I do this and call jax.debug.visualize_array_sharding(mybatch) on the first host, I get the following visualization (assuming 4-way DP and 8-way TP):

┌──────────────────────────────────────────────────────────────────────────────┐
│                            TPU 0,1,2,3,8,9,10,11                             │
├──────────────────────────────────────────────────────────────────────────────┤
│                           TPU 4,5,6,7,12,13,14,15                            │
├──────────────────────────────────────────────────────────────────────────────┤
│                         TPU 16,17,18,19,24,25,26,27                          │
├──────────────────────────────────────────────────────────────────────────────┤
│                         TPU 20,21,22,23,28,29,30,31                          │
└──────────────────────────────────────────────────────────────────────────────┘

To me this signifies that this data batch has now been split across hosts, even though it comes from the first host. Does this mean that data-loading should now only happen on a single host and the batch will then be sharded across the other hosts or am I misunderstanding the above visualization?

Thanks!