A collection of fixes to the mesh dialect
- allow constants in sharding propagation/spmdization
- fixes to tensor replication (e.g. 0d tensors)
- improved canonicalization
- sharding propagation incorrectly generated too many ShardOps
New operation `mesh.GetShardOp` enables exchanging sharding information
(like on function boundaries)
Follow-up to #102598 : as discussed, move tensor sharding implementation
into separate tensor extension lib.
@sogartar @yaochengji, could you take a look at this PR?