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));

  }
};

void TestBitcastOp(Tensor* input_tensor, DataType out_type,
                   TensorShape expected_shape, error::Code expected_code) {
  Status status;
  NodeDef def;
  def.set_op("Bitcast");
  def.set_device(DEVICE_CPU);

  AttrValue typeAttr;
  SetAttrValue(input_tensor->dtype(), &typeAttr);

  AttrValue outTypeAttr;
  SetAttrValue(out_type, &outTypeAttr);

  (*def.mutable_attr())["T"] = typeAttr;

                         error::Code expected_code) {
  // Initialize node used to fetch OpKernel
  Status status;
  NodeDef def;
  def.set_op("ScalarSummary");

  def.set_device(DEVICE_CPU);

  AttrValue valuesTypeAttr;
  SetAttrValue(values->dtype(), &valuesTypeAttr);
  (*def.mutable_attr())["T"] = valuesTypeAttr;

  def.add_input(strings::StrCat("input1: ", DataTypeString(tags->dtype())));

    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_);

                 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));

// `send_device`, `recv_device`, and `send_device_incarnation` set.
TFE_Op* SendOp(TFE_Context* ctx, TFE_TensorHandle* in,
               const std::string& op_name, const std::string& send_device,
               const std::string& recv_device,
               tensorflow::uint64 send_device_incarnation);

// Return a RecvOp op `op_name` with the attributes `send_device`,