// CHECK-DAG: TPUReplicateMetadata
  // CHECK-DAG: TPUCompilationResult
  // CHECK-DAG: TPUReplicatedOutput
  // CHECK: return

  // core_tput should have TPU ops including both replicated inputs and outputs
  // CHECK: func.func private @_func_core_tpu
  // CHECK-DAG: TPUReplicatedInput
  // CHECK-DAG: TPUReplicateMetadata
  // CHECK-DAG: TPUCompilationResult
  // CHECK-DAG: TPUReplicatedOutput
  // CHECK: return

  tf_executor.graph {
    %0 = tf_executor.ControlTrigger {}
    // CHECK: "tf.NoOp"
    %1 = tf_executor.island wraps "tf.NoOp"() : () -> ()
    // CHECK: "tf.TPUCompilationResult"
    %2, %3 = tf_executor.island(%1) wraps "tf.TPUCompilationResult"() : () -> tensor<!tf_type.string>
    tf_executor.fetch %0 : !tf_executor.control
  }
  func.return
}

// CHECK-LABEL: func @preserve_unreachable_tpu_replicated_input

    "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> ()
    // expected-error @+1 {{'tf.TPUCompilationResult' op is not part of the replication region "repl_info" vs "wrong_repl_info"}}
    %1 = "tf.TPUCompilationResult"() {_tpu_compilation_status = "wrong_repl_info"} : () -> tensor<!tf_type.string>
    return %arg0 : tensor<i1>
  }

func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>

    %outputs_5, %control_6 = tf_executor.island(%control_4) wraps "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>

  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "num_replicas"
    value {
      i: 1
    }
  }
}
node {
  name: "TPUCompilationResult"
  op: "TPUCompilationResult"
  input: "^TPUReplicateMetadata"
  attr {
    key: "_tpu_compilation_status"
    value {
      s: "cluster"
    }
  }
}
node {
  name: "output_0_shard_0"

  %1 = "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
  %2 = "tf.Transpose"(%0, %cst_0) {device = ""} : (tensor<1x3x4x3xbf16>, tensor<4xi32>) -> tensor<1x3x3x4xbf16>

    %outputs_2, %control_2 = tf_executor.island wraps "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster"} : () -> tensor<!tf_type.string>
  }
  func.return
}

// -----

// CHECK-LABEL: func @tpu_compilation_status_different_clusters
// CHECK: tf_executor.island
// CHECK: "tf.Const"
// CHECK: tf_executor.island
// CHECK: "tf.TPUCompilationResult"
func.func @tpu_compilation_status_different_clusters() {

  // all graphs should have control dependencies to enforce this but this is
  // currently not the case (see b/177478741).
  const llvm::SmallVector<std::string, 4> ops_to_preserve = {
      "tf.TPUReplicateMetadata", "tf.TPUCompilationResult",
      "tf.TPUReplicatedOutput"};
  bool strict_clusters =
      tensorflow::GetMlirCommonFlags()->tf_mlir_enable_strict_clusters;
  pm.addNestedPass<FuncOp>(

                               llvm::SetVector<Operation*>& merged_set,
                               OpType& found_op) {
  // Find the TPUReplicationMetadata or TPUCompilationResult ops which will be
  // cloned/inserted into each region. We add them to the merged_set so that
  // they're ignored when extracting the four main functions.
  found_op = nullptr;