PR #10495: Enforce same input/output layout for offloading ops
PR #10495: Enforce same input/output layout for offloading ops

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

This patch makes sure that the host offloader will not introduce layout mismatches
after removing the offloading custom calls by constraining the layout assignment
to assign the same layout for the custom call's input and output.
Copybara import of the project:

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

Enforce same input/output layout for offloading ops

Merging this change closes #10495

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10495 from jaro-sevcik:host-offload-call-layout-fix ce6aeac3ac5445ab014d36c92a181dbe7afc35e7
