// name_in is copied here because name may be modified below if
  // rewrite_subgraph_fn is true.
  string name = name_in;
  call_node_def_.set_op(name);
  call_node_def_.set_name(name);
  call_node_def_.set_device(device_);

  if (rewrite_subgraph_fn) {
    std::vector<OutputTensor> arg_source_tensors(args_by_src_.size());
    for (const auto& arg : args_by_src_) {
      arg_source_tensors.at(arg.second) = arg.first;
    }

    if (otherDevice->device_) {
      return SetDevice(otherDevice->device_);
    }
    return ChangeResult::NoChange;
  }

  void print(raw_ostream &os) const override {
    if (device_) {
      os << "<" << device_ << ">";
    } else {
      os << "<no device>";
    }
  }

  ChangeResult SetDevice(mlir::StringAttr device) {
    bool changed = (device != device_);

        const std::string& src_device = ctx->device()->name();
        std::string task_prefix;
        TF_RETURN_IF_ERROR(GetTaskName(src_device, &task_prefix));
        const std::string dst_device =
            absl::StrCat(task_prefix, "/device:CPU:0");
        const std::string& rendezvous_key =
            Rendezvous::CreateKey(src_device, /*src_incarnation=*/1, dst_device,

  se_.create_event = [](const SP_Device* const device, SP_Event* event,
                        TF_Status* const status) -> void {
    *event = new SP_Event_st(123);
  };
  se_.destroy_event = [](const SP_Device* const device,
                         SP_Event event) -> void { delete event; };
  se_.get_event_status = [](const SP_Device* const device,

  // Used for printing. Must be null-terminated.
  const char* pci_bus_id;
} SP_Device;

#define SP_DEVICE_STRUCT_SIZE TF_OFFSET_OF_END(SP_Device, pci_bus_id)

typedef struct SE_CreateDeviceParams {
  size_t struct_size;
  void* ext;        // reserved for future use
  int32_t ordinal;  // device index

  SP_Device* device;  // Input/output, struct_size set by TF for plugin to read.

                 SP_Event event, TF_Status* const status) {}
void WaitForEvent(const SP_Device* const device, SP_Stream stream,
                  SP_Event event, TF_Status* const status) {}
void CreateTimer(const SP_Device* const device, SP_Timer* timer,
                 TF_Status* const status) {}
void DestroyTimer(const SP_Device* const device, SP_Timer timer) {}

  // CHECK-NEXT: %[[RESULT_3:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_2]], %[[RESULT_2]])
  %3 = corert.executeop(%0) "tf.AddV2"(%2, %2) {T = f32, device = "/device:CPU:0"} : 1
  // CHECK-NEXT: %[[RESULT_4:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
  // CHECK-NEXT: %[[RESULT_5:.*]] = corert.get_dst_tensor_type %[[RESULT_3]], %[[RESULT_4]]
  // CHECK-NEXT: %[[RESULT_6:.*]] = corert.transfer %[[RESULT_3]], %[[RESULT_4]], %[[RESULT_5]]

#include "tensorflow/core/framework/tensor_testutil.h"
#include "tsl/lib/core/status_test_util.h"

namespace tensorflow {
namespace {

class XlaHostSendRecvDeviceContextTest : public ::testing::Test {
 public:
  void SetDevice(const string& device_type) {
    auto device_factory = DeviceFactory::GetFactory(device_type);
    SessionOptions options;
    std::vector<std::unique_ptr<Device>> devices;
    Status s = device_factory->CreateDevices(

  tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  return true;
}();

class DeviceContextTest : public ::testing::Test {
 public:
  void SetDevice(const string& device_type) {
    auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
    rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type));

    auto it = value_to_device.find(operand);
    if (it == value_to_device.end()) return llvm::StringRef();

    if (new_device.empty()) {
      new_device = it->getSecond();
      continue;
    }

    if (new_device != it->getSecond()) return llvm::StringRef();
  }

  return new_device;
}

// Propagates devices from function arguments.
void PropagateDevicesFromArguments(
    func::FuncOp func,