PR #11354: Passthrough sharding propagation for host offloading ops
PR #11354: Passthrough sharding propagation for host offloading ops

Imported from GitHub PR https://github.com/openxla/xla/pull/11354

To support host memory offloading we propagate sharding on the offloading operations (MoveToDevice and MoveToHost custom calls) so that spmd partitioner does not insert collectives at wrong places.
Copybara import of the project:

--
6a4d73c56e618e921fa32e5e03756e40986a8ad3 by Jaroslav Sevcik <jsevcik@nvidia.com>:

Passthrough sharding propagation for host offloading ops

Merging this change closes #11354

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11354 from jaro-sevcik:host-offloading-sharding-propagation 6a4d73c56e618e921fa32e5e03756e40986a8ad3
