%embedding_pass_trigger = "tf.Const"() {_embedding_pipelining = "forward", _replication_info = "repl_info", value = dense<2> : tensor<i32>} : () -> tensor<i32>
    "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> ()
    "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> ()

// CHECK-NEXT: = "tf.AddV2"
// CHECK-SAME: _replication_info
// CHECK: tf_executor.island wraps "tf.AddV2"
    %outputs, %control = tf_executor.island wraps "tf.Const"() {_xla_compile_device_type = "TPU", _replication_info = "cluster", value = dense<1> : tensor<i32>} : () -> tensor<i32>

    %embedding_pass_trigger = "tf.Const"() {_embedding_pipelining = "forward", _replication_info = "repl_info", value = dense<2> : tensor<i32>} : () -> tensor<i32>
    "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> ()
    "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> ()

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

  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'}

    For example, a non replicated `tf_device.cluster_func`:

    ```mlir
    func @tf_tpu_rewrite(%arg0: tensor<i8>) {
      %0 = "tf_device.cluster_func"(%arg0) {_xla_compile_device_type = "TPU", _replication_info = "cluster0", func = @func} : (tensor<i8>) -> tensor<i8>
      return
    }
    ```

    will be rewritten as:

    ```mlir
    func @tf_tpu_rewrite(%arg0: tensor<i8>) {

    %6 = "tf.ReadVariableOp"(%arg9) : (tensor<!tf_type.resource<tensor<i64>>>) -> tensor<i64>

std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUPartitionedOpConversionPass();

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass();

// Creates a pass that cleans up `_replication_info` attribute on operations
// that are inside a cluster.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUClusterCleanupAttributesPass();

// Creates a pass that removes Identity/IdentityN ops from a cluster.

			bui.ReplicaSize = flat.ReplicationStats.ReplicaSize
			bui.ReplicaCount = flat.ReplicationStats.ReplicaCount

			bui.ReplicationInfo = make(map[string]BucketTargetUsageInfo, len(flat.ReplicationStats.Targets))
			for arn, stat := range flat.ReplicationStats.Targets {
				bui.ReplicationInfo[arn] = BucketTargetUsageInfo{
					ReplicationPendingSize:  stat.PendingSize,
					ReplicatedSize:          stat.ReplicatedSize,