func.func @multiple_replicated_interleaved(%arg0: !tf_res) {
  "tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "a", num_replicas = 2, topology = "topology"} : () -> ()
  "tf.TPUReplicateMetadata"() {_xla_compile_device_type = "TPU", _replication_info = "b", num_replicas = 2, topology = "topology"} : () -> ()

  // CHECK-NOT: TPUReplicatedInput
  // 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

func.func @num_replicas_replicated(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  %0:2 = tf_executor.graph {
    %control = tf_executor.island() wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", device = "/device:TPU:0", num_replicas = 2, topology = "topology"} : () -> ()

// CHECK-LABEL: func nested @_tpu_v1_compat_outlined_func0(%arg0: tensor<i1>) -> tensor<i1>
// CHECK-NEXT: tf.TPUReplicateMetadata
// CHECK-NEXT: tf.opA

// CHECK-LABEL: func nested @_tpu_v1_compat_outlined_func1(%arg0: tensor<i1>, %arg1: tensor<f32>) -> (tensor<i1>, tensor<i32>)
// CHECK-NEXT: tf.TPUReplicateMetadata
// CHECK-NEXT: tf.opA
// CHECK-NEXT: tf.opA
// CHECK-NEXT: tf.SomeOp

// -----

    %control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], name = "TPUReplicateMetadata", num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, step_marker_location = "STEP_MARK_AT_ENTRY", topology...

      i: 1
    }
  }
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "cluster/pivot"
  op: "NoOp"
}
node {
  name: "TPUReplicateMetadata"
  op: "TPUReplicateMetadata"
  input: "^cluster/pivot"
  attr {
    key: "_tpu_replicate"
    value {
      s: "cluster"
    }
  }
  attr {
    key: "computation_shape"
    value {
      list {
      }

  // CHECK: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])
  // CHECK-SAME: _XlaSharding = "123"
  // CHECK-SAME: partition_dims = [2, 1]
  // CHECK: [[RI:%.*]] = "tf.TPUReplicatedInput"([[PI_0]])
  "tf.TPUReplicateMetadata"() {num_cores_per_replica = 2 : i64, num_replicas = 1 : i64} : () -> ()
  %pi_0 = "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "123", partition_dim = 0 : i64} : (tensor<10x3xf32>, tensor<10x3xf32>) -> tensor<20x3xf32>

    return
  }
  // expected-error @+1 {{number of tf.TPUReplicateMetadata in loop body is not 1}}
  func.func private @while_body(%arg0: tensor<i1>) -> (tensor<i1>) {
    // metadata ops
    %embedding_pass_trigger = "tf.Const"() {_embedding_pipelining = "forward", _replication_info = "repl_info", value = dense<2> : tensor<i32>} : () -> tensor<i32>

    %1 = tf_executor.island wraps "tf.NoOp"() : () -> ()
    // CHECK: "tf.TPUReplicateMetadata"
    %2 = tf_executor.island(%1) wraps "tf.TPUReplicateMetadata"() {allow_soft_placement = false, computation_shape = [], device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()

  // CHECK: "tf.TPUReplicateMetadata"()
  // CHECK: [[RI_0:%.*]] = "tf.TPUReplicatedInput"([[ARG0]])
  // CHECK-SAME: is_packed = true
  // CHECK: [[RI_1:%.*]] = "tf.TPUReplicatedInput"([[ARG0]])
  // CHECK-SAME: is_packed = true
  // CHECK: [[PI:%.*]] = "tf.TPUPartitionedInputV2"([[RI_0]], [[RI_1]])
  // CHECK-SAME: is_packed = false