xla

0a6e1659 - Set replicated output false (#5149)

Set replicated output false (#5149) * [SPMD] index local shards correctly in the execution arguments for replication. * [SPMD] support partial replication with xs.ShardingSpec * [SPMD] output sharding by default & update ShardingType commnet