// CHECK-LABEL: func @replicate_with_devices
func.func @replicate_with_devices() {
  tf_device.replicate() {n = 2 : i32, devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}} {
    tf_device.return
  }
  func.return

// CHECK:      tf_device.replicate
// CHECK-SAME: devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}
// CHECK-SAME: n = 2
// CHECK-NEXT:   tf_device.return
}

// -----

    // CHECK: %[[REPLICATE:[0-9]*]]:2 = tf_device.replicate
    // CHECK-SAME: ([%[[A_OUTPUT]], %[[ARG_0]]] as %[[RI_0:[a-z0-9]*]]: tensor<?xi32>)
    // CHECK-SAME: devices = {TPU_REPLICATED_CORE_0 = ["/job:worker/replica:0/task:0/device:TPU:0", "/job:worker/replica:0/task:0/device:TPU:1"], TPU_REPLICATED_HOST_0 = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:CPU:0"]}

    %cst_1 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
    %cst_2 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

      builder.create<mlir::TF::XlaSendToHostOp>(while_op.getLoc(), condition,
                                                condition_send_recv_key);
      // device_ordinal0 is the ordinal of TPU_REPLICATED_CORE_0 and is only
      // used in the replicated case.
      Value device_ordinal0 = nullptr;
      if (!core_to_device_ordinal.empty())
        device_ordinal0 = core_to_device_ordinal[0];