xla

aac03da0 - 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