jax

d621737f - [XLA:Collective] Expose a factory for constructing HLOSharding with explicit device ordering.