Why the order of visualize_array_sharding() is so strange?

[yixiaoer]

The order of devices in the output of visualize_array_sharding() is very strange. It is in the order of 0, 1, 2, 3, 6, 7, 4, 5 instead of 0, 1, 2, 3, 4, 5, 6, 7. Why is that?

[jakevdp]

Hi - thanks for the question! In auto-sharding contexts, JAX does its best to lay out data in a manner that is most efficient given the topology of the device(s). In the case of TPU, the topology may be a 3D mesh or a torus, depending on the system (look for Interconnect Topology in the Cloud TPU Docs).

You can get a representation of the topology you are working with by printing the coordinates of the TPU devices:

import jax
print([d.coords for d in jax.devices()])

Does that answer your question?

[yixiaoer]

Yes, it makes sense. Thanks for the explanation!