if (cc_ctx->op_device_context() == nullptr) {  // CPU Device
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is not supported for a CPU device.");
    return nullptr;
  } else if (!cc_ctx->op_device_context()->IsPluggableDevice()) {
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is only supported for pluggable devices.");
    return nullptr;

FunctionDef SquareFunc() {
  return FunctionDefHelper::Define(
      // Function Name
      "SquareFunc",
      // Args
      {"x: float"},
      // Returns
      {"y: float"},
      // Attr def
      {},
      // Nodes
      {{/*ret=*/{"y"},
        /*op=*/"Square",
        /*arg=*/{"x"},
        /*attr=*/{{"T", DT_FLOAT}},
        /*dep=*/{},
        /*device=*/"",
        /*name=*/"square"}});
}

inline constexpr absl::string_view kDeviceAssignmentAttr = "device_assignment";

// A TPU device for execution alongside its associated host CPU device.
struct TPUDeviceAndHost {
  TPUDeviceAndHost() = default;
  TPUDeviceAndHost(llvm::StringRef device, llvm::StringRef host)
      : device(device), host(host) {}

  std::string device;
  std::string host;
};

    value_to_device.insert({result, device});
  }
}

// Checks if an operation can have TPU devices propagated through.
bool IsSupportedOpToSetDevice(Operation& op) {
  return IsSupportedExecutorOp(op) ||
         isa<TF::IdentityOp, TF::IdentityNOp, TF::ShapeOp>(op);
}

// Finds nonconflicting TPU device for an operation from its operands. If an
// operand has no device or a non TPU device, or if there are conflicting

                            void* device_info) = nullptr;

  // Pins the op to `device` based on inputs to `op`. Returns true
  // signifying to pin to the current custom device. Returns false
  // to pin to the physical device.
  //
  // This function is guaranteed to be called only when all of the custom-device
  // inputs are on this device.
  bool (*shall_pin_to_this_device)(const TFE_Op* op, TF_Status* s) = nullptr;

  virtual void SetExecutorForThread(EagerExecutor* executor) = 0;

  // Return a list of local tensorflow::Device*.
  // TODO(tfrt-devs): We shouldn't expose legacy device in this API.
  virtual std::vector<tensorflow::Device*> ListLocalTfDevices() = 0;

  // Return a list of all tensorflow::Device*.
  virtual std::vector<tensorflow::Device*> ListAllTfDevices() = 0;

  //===--------------------------------------------------------------------===//

  auto node_device = [&](Operation *n) -> StringAttr {
    auto device = n->getAttrOfType<StringAttr>(device_id);
    if (!device || device.getValue().empty()) return caller_device;

    DeviceNameUtils::ParsedName ndef_parsed_device;
    if (!DeviceNameUtils::ParseFullName(device.getValue().str(),
                                        &ndef_parsed_device))
      return device;
    DeviceNameUtils::MergeUnsetDevNames(&ndef_parsed_device,

    // The name of the compilation device (e.g., "XLA_CPU_JIT");
    string compilation_device_name;

    // If 'use_multiple_streams' is true, we create separate streams for
    // compute, host-to-device, and device-to-host communication.
    bool use_multiple_streams = false;

    // If true, the XLA devices with the same device ordinal will share the same
    // compute stream. Otherwise each XLA device will having their own compute

                {n = 2 : i32,
                 devices = {DEVICE_ALIAS_0 = ["/DEVICE:0", "/DEVICE:1"],
                            DEVICE_ALIAS_1 = ["/DEVICE:2", "/DEVICE:3"]}} {
  // Inside the region, %0, %2, %4, and %6 corresponds to
  // "/DEVICE:0"/"/DEVICE:2" and %1, %3, %5, and %7 corresponds to
  // "/DEVICE:1"/"/DEVICE:3", depending on which device alias is used.
  %k = "tf_device.launch"() ( {

    args: a list of fallback tensors.
    device: the name of the tensorflow device that is associated with the
      input fallback tensors.

    results: the converted corert tensorhandles.
  }];

  let arguments = (ins
    Variadic<TFTensorType>:$args,
    StrAttr:$device
  );

  let results = (outs
    Variadic<CoreRT_TensorHandleType>:$results
  );