// Returns pointer to the variable from 'session' that 'var_handle_op'
// refers to which is in 'device_name' device. If failed to fetch the value null
// will be returned.
// Note, caller is responsible for Unref the variable.
tensorflow::Var* GetVariableFromSession(mlir::TF::VarHandleOp var_handle_op,
                                        llvm::StringRef device_name,
                                        const tensorflow::DeviceMgr* mgr);

  const char* second_device = "/job:worker/replica:0/task:2/device:CPU:0";
  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
  std::array<const char*, 2> underlying_devices{first_device, second_device};
  RegisterParallelDevice(context.get(), device_name, underlying_devices,
                         status.get());

}

tensorflow::Var* GetVariableFromSession(mlir::TF::VarHandleOp var_handle_op,
                                        llvm::StringRef device_name,
                                        const tensorflow::DeviceMgr* mgr) {
  tensorflow::Device* device = nullptr;
  if (!mgr || !mgr->LookupDevice(StringRefToView(device_name), &device).ok())
    return nullptr;
  tensorflow::Var* var_ptr = nullptr;
  const auto& container = var_handle_op.getContainer().str();

  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
  std::array<const char*, 2> underlying_devices{first_device, second_device};
  RegisterParallelDevice(context, device_name, underlying_devices,
                         status.get());
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

                                                 const string& device_name,
                                                 int device_ordinal) {
  return Device::BuildDeviceAttributes(
      absl::StrCat(name_prefix, "/device:", device_name, ":", device_ordinal),
      DeviceType(device_name), Bytes(16ULL << 30), DeviceLocality(),
      absl::StrCat("device: ", device_name, " device"));
}

}  // namespace

    OwnedTFStatus c_status(TF_NewStatus());
    SP_Stream stream_handle = static_cast<CStream*>(stream)->Handle();
    SP_DeviceMemoryBase device_mem = DeviceMemoryBaseToC(location);
    stream_executor_->mem_zero(&device_, stream_handle, &device_mem, size,
                               c_status.get());
    return StatusFromTF_Status(c_status.get());
  }

}

TF_StringView TF_GetDeviceName(TF_OpKernelContext* ctx) {
  const auto& device_name =
      reinterpret_cast<tensorflow::OpKernelContext*>(ctx)->device()->name();
  TF_StringView device_name_sv;
  device_name_sv.data = device_name.data();
  device_name_sv.len = device_name.length();
  return device_name_sv;
}

#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)

    const char* device_name = TFE_OpGetDevice(matmul, status);
    ASSERT_TRUE(strstr(device_name, "GPU:0") != nullptr);

    TFE_OpSetDevice(matmul, "CPU:0", status);
    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
    device_name = TFE_OpGetDevice(matmul, status);
    ASSERT_TRUE(strstr(device_name, "CPU:0") != nullptr);
  }

  TFE_DeleteOp(matmul);

class DeviceSetup {
 public:
  void AddDevicesAndSetUp(
      const std::vector<std::string>& device_names,
      const std::optional<FunctionDef>& fdef = std::nullopt);
  Device* GetDevice(const string& device_name);
  FunctionLibraryRuntime* flr() { return flr_; }

 private:
  FunctionLibraryRuntime* flr_;
  std::unique_ptr<DeviceMgr> device_mgr_;
  std::unique_ptr<FunctionLibraryDefinition> lib_def_;

  //===--------------------------------------------------------------------===//
  virtual CustomDeviceOpHandler& GetCustomDeviceOpHandler() = 0;

  // Returns whether `device_name` is registered as a custom device.
  virtual bool IsCustomDevice(const string& device_name) = 0;

  // Register a custom device. It will return error is the device name is
  // already registered.