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

      }) {is_stateless = true, _xla_compile_device_type = "TPU", _replication_info = "x" } : (tensor<i1>) -> (tensor<f32>)
    tf_device.return %3 : tensor<f32>
  // CHECK: {_replication_info = "x", _xla_compile_device_type = "TPU", cluster_attr = "cluster_attr", device = "y"}
  }) {cluster_attr = "cluster_attr", _xla_compile_device_type = "TPU", _replication_info = "x", device = "y"} : () -> tensor<f32>
  // CHECK: "tf.Add"

  // 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>) -> ()

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s
// RUN: tf-opt %s -tf-replicated-clustering-bridge-v1 -tfrt-lower-cluster-to-runtime-ops-tpu -tf-dialect-to-executor-v1 2>&1 | FileCheck %s

// This test verifies there is no warning about shape inference failure in TPU
// bridge in handling multiple usage of the same function.

// Since it is possible that this warning may become an error in the future,

  );

  let assemblyFormat = "attr-dict";
}

def CompileAndExecuteOp : TensorflowMlrtTpu_Op<"compile_and_execute"> {
  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.

// RUN: tf-opt %s -tf-replicated-clustering-bridge-v2 -tfrt-lower-cluster-to-runtime-ops-tpu 2>&1 | FileCheck %s

// TPUReshardVariables should be inserted even when While functions' shapes are
// different than While operand shapes. Test the whole tf-tpu-bridge because
// correct insertion of TPUReshardVariables depends on multiple passes including
// TPUVariableRuntimeReformatting, ShapeInference, WhileRegion canonicalization,
// and TPUMergeVariablesWithExecute.

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-tpu-annotate-dynamic-shape-inputs %s | FileCheck %s

// Test that annotate the inputs of the cluster func to be dynamic shaped.

module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
   func.func @main(
      %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},

  }];
}

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.

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace mlir {
namespace TFTPU {

// Creates a pass that rewrites `tf_device.launch_func` on TPUs into TPU runtime
// ops.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTPURewritePass(
    llvm::StringRef module_name = llvm::StringRef());

// Creates a pass that adds ops which perform formatting on variables at