%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>

    // expected-error @+1 {{'tf.Identity' op is expected to have either both or none of '_replication_info' and '_xla_compile_device_type' attributes}}

  // CHECK-NEXT: [[READ:%.*]] = "tf.ReadVariableOp"([[ARG2]])
  // CHECK-NEXT: [[CLUSTER:%.*]]:2 = "tf_device.cluster_func"([[ARG0]], [[ARG1]], [[READ]])
  // 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>)

      "tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", num_relicas = 1, num_cores_per_replica = 1, topology = "<topology>", device_assignment = [], padding_map = []} : () -> ()
      %replicated_input = "tf.TPUReplicatedInput"(%arg0) : (tensor<i32>) -> tensor<i32>
      %identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info = "cluster"} : (tensor<i32>) -> tensor<i32>

  func.func @missing_compilation_and_replication_attributes() {
    // expected-error@+1 {{'tf_device.cluster_func' op is expected to have either both or none of '_replication_info' and '_xla_compile_device_type' attributes}}

  return success();
}

// Forms clusters with ops of the same `_replication_info` attribute under a
// block.
//
// For a given block, clusters are formed via grouping ops by
// `_replication_info` attributes. For every cluster formed:
//   1. Find associated TPUReplicateMetadata attributes with the same
//      `_replication_info` attribute.
//   2. Find users not in cluster that are interleaved between cluster ops.

        // already have a cluster assigned to them (via `_replication_info`
        // attribute value).
        // `tf.Identity` op is also treated as special tpu ops since it can play
        // a role as connection between `tf.TPUReplicatedInput` or
        // `tf.TPUPartitionedInput`. For example, we have the follow pseudocode:
        // %0 = tf_executor.island wraps "tf.OpA" (){_replication_info = 'c'}

// Verifies that all MLIR Ops have the expected attributes.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyClusteringPass();

// Creates a pass that forms clusters from operations of the same
// `_replication_info` attribute.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUClusterFormationPass(bool strict_clusters = false);

// Creates a pass that extracts outside compilation (Host ops inside device

    "_cluster_outlined_function_name";
// Marks a node for replication. The attribute value indicates the replication
// metadata op.
inline constexpr llvm::StringRef kReplicationInfoAttr = "_replication_info";
// Marks a node for XLA-TPU compilation. The attribute value indicates the
// associated compilation cluster and replication metadata op.
inline constexpr llvm::StringRef kTpuReplicateAttr = "_tpu_replicate";

    ```

    wll be replaced by `_replication_info="cluster"` and  `_xla_compile_device_type="TPU"`.

    ```mlir