// PjRtTensorBuffer. The PjRtTensorBuffer holds the pointer to the device
  // memory. It also owns the PjRtBuffer. If output_tensor does not use
  // PjRtTensorBuffer and the opaque device memory is the same, update the
  // output_tensor->buf_ so that the same device memory will not be double-free.
  // Otherwise a new Tensor will be created with the PjRtTensorBuffer.
  //
  // TODO(b/289001822): Create a unit test to cover this function.

  }

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:
  se::Stream* stream_;                        // Not owned.

// after the _TPUCompileMlir op.
std::unique_ptr<OperationPass<func::FuncOp>> CreateEmbeddingProgramKeyPass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_EMBEDDINGSEQUENCINGPASS
#define GEN_PASS_DECL_EMBEDDINGPIPELININGPASS
#define GEN_PASS_DECL_EMBEDDINGPROGRAMKEYPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

}  // namespace TFDevice

// Creates a pass that merges device variable reads/updates into the surrounded
// TPUExecute node. This allows the execute node to perform in-place variable
// updates.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUMergeVariablesWithExecutePass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_TPUMERGEVARIABLESWITHEXECUTEPASS
#define GEN_PASS_DECL_TPUREWRITEPASS
#define GEN_PASS_DECL_TPUVARIABLERUNTIMEREFORMATTINGPASS

// available. It's completely valid to have a device without any metadata
// attached to it.
class RuntimeDevices {
  using DeviceNameUtils = ::tensorflow::DeviceNameUtils;
  using ParsedName = ::tensorflow::DeviceNameUtils::ParsedName;

 public:
  // Adds a device with and empty metadata. Device can be of any type.
  void AddDevice(const ParsedName& device);

  // Adds a GPU device with GPU specific metadata.

  }];
}

// This corresponds to tensorflow::Device* .
def TFDeviceType : DialectType<TensorflowMlrt_Dialect,
    CPred<"$_self.isa<::tensorflow::tf_mlrt::TFDeviceType>()">, "!tf_mlrt.device type">,
    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFDeviceType>()"> {
  let description = [{
    `!tf_mlrt.device type` represents a tensorflow::device.
  }];
}

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_OPS_DEVICE_HELPER_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_OPS_DEVICE_HELPER_H_

namespace mlir {

class Operation;

namespace TF {

class RuntimeDevices;

// Returns true if at least one GPU device is available at runtime.
bool CanUseGpuDevice(const RuntimeDevices &devices);

// Classes and utilities that work with StreamExecutor C API for internal use.
// This includes functions used for device registration and interfaces needed
// for testing.
#ifndef TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_
#define TENSORFLOW_C_EXPERIMENTAL_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_

#include "tensorflow/c/experimental/stream_executor/stream_executor.h"

// and reachable in the call graph from the cluster.
// For each op walked, `callback` is applied to the op, the root cluster, and
// the root cluster's host device. `callback` returning WasInterrupted
// indicates failure.
// The host device is null when the tpu_cluster HasModelParallelism: The
// HasModelParallelism case is currently unsupported in combination with
// outside compilation.

namespace {

constexpr char kDeviceAttr[] = "device";

#define GEN_PASS_DEF_DEVICEATTRIBUTETOLAUNCHPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_device_passes.h.inc"

struct DeviceAttributeToLaunch
    : public impl::DeviceAttributeToLaunchPassBase<DeviceAttributeToLaunch> {
  void runOnOperation() override;
};

void WrapOpInLaunch(Operation* op, llvm::StringRef device) {
  OpBuilder builder(op);