xla

6852df1d - SPMD cross slice-replication using partial_replication sharding (#5411)

SPMD cross slice-replication using partial_replication sharding (#5411) * Revert "Support unordered sharding spec for partial replication (#5316)" * Update test_2d_tensor_3d_mesh unit test to surface a bug * Use partial replication for 2D tensor over 3D mesh sharding