xla

78c6840c - Support unordered sharding spec for partial replication (#5316)

Support unordered sharding spec for partial replication (#5316) * Suport unordered sharding spec for partial replication * add 4d test * handle 2d tensor with 2d mesh case * refactoring