%1 = "tf.Relu"(%0) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
    %2 = "tf.Add"(%0, %1) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %3 = "tf.Relu"(%arg1) {device = "tpu1"} : (tensor<?xi32>) -> tensor<?xi32>
    %4 = "tf.Add"(%3, %arg1) {device = "tpu1"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %5 = "tf.Relu"(%4) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>

  const std::string cpu0 = "/job:worker/replica:0/task:0/device:CPU:0";
  const std::string gpu0 = "/job:worker/replica:1/task:2/device:GPU:0";
  const std::string gpu1 = "/job:worker/replica:1/task:2/device:GPU:1";

  llvm::SmallVector<std::unique_ptr<Device>, 2> devices;
  devices.push_back(FakeDevice::Make(cpu0));
  devices.push_back(FakeDevice::Make(gpu0, "compute capability: 7.0"));

  auto status = tensorflow::DeviceNameUtils::DeviceNameToCpuDeviceName(
      tpu0_device, cpu0_device);
  if (!status.ok())
    return op->emitError()
           << "error in converting TPU0 to CPU0. The TPU device is "
           << tpu0_device;
  return mlir::success();
}

}

def TPUPartitionedOpConversionPass : Pass<"tf-tpu-partitioned-op-conversion", "mlir::func::FuncOp"> {
  let summary = "Rewrite all TPU Partitioned ops into their V2 counterparts.";

  let constructor = "TFTPU::CreateTPUPartitionedOpConversionPass()";
}

def TPUValidateInputsPass : Pass<"tf-tpu-validate-inputs", "ModuleOp"> {
  let summary = "Validates inputs to the TPU TF/XLA bridge";

  let description = [{

#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/status_helper.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/stream_executor/tpu/tpu_node_context.h"
#include "xla/stream_executor/tpu/tpu_platform.h"
#include "xla/stream_executor/tpu/tpu_platform_interface.h"
#include "xla/stream_executor/tpu/tpu_stream_interface.h"

  // CHECK-SAME: _replication_info = "write", _xla_compile_device_type = "TPU"
  %0:2 = "tf_device.cluster_func"(%arg0, %arg1) {_replication_info = "write", _xla_compile_device_type = "TPU", func = @write_func} : (tensor<i32>, tensor<f32>) -> (tensor<f32>, tensor<i32>)
  // CHECK-NEXT: "tf.AssignVariableOp"([[ARG2]], [[CLUSTER]]#1)
  "tf.AssignVariableOp"(%arg2, %0#1) : (tensor<*x!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()

  }];
}

def TFTPUCompileAndExecuteOp : TensorflowMlrt_Op<"tf_tpu_compile_and_execute", [TF_MustExecute]> {
  let summary = "tpu compile and execute operation";
  let description = [{
    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.

    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.

  }];

  let constructor = "tensorflow::tf2xla::internal::CreateVerifyClusteringPass()";
}

def TPUClusterFormationPass : Pass<"tf-tpu-cluster-formation", "ModuleOp"> {
  let summary = "Forms clusters from operations assigned to the same TPU computation";

  let description = [{
    TPU computations from the frontend are composed of a `tf.TPUReplicateMetadata`
    op, a subgraph of ops (TensorFlow Dialect) each with a matching

module attributes {tf.versions = {producer = 888 : i32}} {
  // CHECK-LABEL: func @single_gpu_cluster_func
  func.func @single_gpu_cluster_func(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

  %3 = "tf.ReadVariableOp"(%0) {_xla_compile_device_type = "TPU", _replication_info = "a"} : (!tf_res) -> tensor<f32>
  %4 = "tf.ReadVariableOp"(%1) {_xla_compile_device_type = "TPU", _replication_info = "b"} : (!tf_res) -> tensor<f32>
  %5 = "tf.ReadVariableOp"(%2) {_xla_compile_device_type = "TPU", _replication_info = "c"} : (!tf_res) -> tensor<f32>
  %6 = "tf.Identity"(%3) {_xla_compile_device_type = "TPU", _replication_info = "a"} : (tensor<f32>) -> tensor<f32>