xs.make_mesh

🚀 Feature

We should provide a xs.make_mesh method similar to JAX: https://docs.jax.dev/en/latest/_autosummary/jax.make_mesh.html

Motivation

The ordering of devices affects the performance of collective operations:

• Weights are sharded differently and the layout of tensors after an all-gather is different.
• The performance of a collective op itself depends on what network connections it uses.

Pitch

The xs.make_mesh method should attempt to create a reasonably optimized device ordering given a mesh shape referencing how JAX does it.

We could also deviate from JAX and add additional options that fit our own needs.

cc @yaochengji