[SPMD] Use UNKNOWN sharding type for implicit replication
This is part of auto-sharding PoC described in #6322 . This PR addresses the following:
- Adapt xla::OpSharding::UNKNOWN for implicit replication type.
- Refactoring of ShardingUtil::GetOutputSharding, ShardingUtil::CreateShardedData
- Remove resolved TODOs
It is important to distinguish implicit and explicit replciation for XLA compiler to generate shardings for.
I guess my question is why do we need to distinguish implicit sharding and explicit sharding?
I guess my question is why do we need to distinguish implicit sharding and explicit sharding?
Good question, updated the description. TLDR, XLA needs to know what is implicit and expcliit so that can consider implicit for auto-sharding.
Can you remind me why do we need to explicit set the output sharding to be replicated before and how is left it to
unknown(implicit sharding) will do the same thing?
For the outputs, we set them to explicit replicated. Let me sync with you offline!
Ok CI passed, will rebase and land