t.scalar<tensorflow::ResourceHandle>()().set_device(kDeviceName);

      outputs->push_back(t);
    } else if (output_name == "var2") {
      Tensor t = Tensor(tensorflow::DT_RESOURCE, tensorflow::TensorShape({1}));
      t.scalar<tensorflow::ResourceHandle>()().set_name("var2");
      t.scalar<tensorflow::ResourceHandle>()().set_device(kDeviceName);

      outputs->push_back(t);

  auto& test_params = GetParam();
  DataType dtype = std::get<0>(test_params);
  TensorShape shape(std::get<1>(test_params));

  SavedVariable saved_variable;
  saved_variable.set_dtype(dtype);
  saved_variable.set_device("/job:localhost/replica:0/task:0/device:CPU:1"),
      shape.AsProto(saved_variable.mutable_shape());

  std::unique_ptr<Variable> var;
  TF_ASSERT_OK(internal::LoadSavedVariable(context(), saved_variable, &var));

    // Target the XLA CPU/GPU backends.
    VLOG(2) << "Replacing with XlaLaunch";
    VLOG(2) << "Device is " << launch->requested_device();
    def.set_op("XlaLaunch");
    def.set_device(launch->requested_device());
    AddNodeAttr("Tconstants", DataTypeVector{}, &def);
    AddNodeAttr("Targs", arg_types, &def);
    AddNodeAttr("Nresources", num_variables, &def);

  // associated NodeDef is not updated.
  for (int i = 0, e = inst->getNumOperands(); i < e; ++i) {
    node_def->add_input();
  }
  if (auto attr = inst->getAttrOfType<mlir::StringAttr>("device")) {
    node_def->set_device(std::string(attr.getValue()));
  }

  // Add the node debug info.
  TF_RETURN_IF_ERROR(ConvertLocation(
      inst->getLoc(), name, node_def->mutable_experimental_debug_info()));

  return node_def;
}

    Var* var = CreateVariable<T>(name, shape, data);
    ResourceMgr* rm = device_->resource_manager();
    TF_ASSERT_OK(rm->Create(rm->default_container(), name, var));

    ResourceHandle handle;
    handle.set_device(device_->name());
    handle.set_container(rm->default_container());
    handle.set_name(name);
    TypeIndex type_index = TypeIndex::Make<Var>();
    handle.set_hash_code(type_index.hash_code());

  // 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;
    }

        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,

#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(

    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,