xla

85a4f141 - [SPMD] Support partial replication. (#4788)

[SPMD] Support partial replication. (#4788) * [SPMD] Support partial replication. * [SPMD] index local shards correctly in the execution arguments for replication. * [SPMD] support partial replication with xs.ShardingSpec