PR #11435: Spmd partitioning for host memory offloading ops
PR #11435: Spmd partitioning for host memory offloading ops

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

This partitions MoveToHost and MoveToDevice custom calls in a way that
tries to avoid resharding on the host. That is, MoveToHost reshards input,
MoveToDevice reshards output.
Copybara import of the project:

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

Partition host offloading ops

Merging this change closes #11435

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11435 from jaro-sevcik:host-offloading-spmd-partitioning a6d01010f3b468de02ebee8a0c59f9e28ea4297d
