// CHECK-SAME: devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]
          %rep:2 = tf_device.replicate([%arg0, %arg1] as %arg30: tensor<*x!tf_type.resource<tensor<f32>>>,
                              [%arg2, %arg3] as %arg31: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>)
                  {_mirrored_variable_indices = [0, 1], devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]}, n = 2 : i32} {

// CHECK:  tf_device.replicate({{\[}}[[VAL_0]], [[VAL_1]]] as [[VAL_4:%.*]]: tensor<*xf32>) {devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]}, n = 2 : i32} {
// CHECK:    [[OP_A:%.*]] = "tf.opA"([[SHAPE]]) {device = "TPU_REPLICATED_CORE_0"} : (tensor<?xi32>) -> tensor<*xi32>
// CHECK:    [[LAUNCH_B:%.*]] = "tf_device.launch"() <{device = "TPU_REPLICATED_CORE_0"}> ({
// CHECK:      [[OP_B:%.*]] = "tf.opB"([[SHAPE]]) : (tensor<?xi32>) -> tensor<*xi32>

    // CHECK: "tf_device.parallel_execute"
    "tf_device.parallel_execute"() ({
      // CHECK: "tf_device.launch"
      // CHECK-SAME: device = "TPU_REPLICATED_CORE_0"
      // CHECK: "tf.TPUExecute"(%[[R0]], %[[COMPILE]]#1)

    }
    ```

    will be rewritten as:

    ```mlir
    func @tf_tpu_rewrite(%arg0: tensor<i8>, %arg1: tensor<i8>) {
      %0:2 = tf_device.replicate([%arg0, %arg1] as %arg2: tensor<i8>) {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"]}, n = 2 : i32} {

  "tf_device.replicate" () ({
// expected-error@-1 {{'tf_device.replicate' op expects number of devices (2) to be equal to 'n' (3)}}
  ^entry:
    tf_device.return
  }) {devices = {TPU_REPLICATED_CORE_0 = ["/DEVICE:0", "/DEVICE:1"]}, n = 3 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
}

// -----

// Check that replicate op's 'devices' attribute must consist of dictionary

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

// CHECK-LABEL: func @device_ordinals
func.func @device_ordinals() {
  tf_executor.graph {
    %0:1 = tf_executor.island {
      tf_device.replicate {n = 2 : i32, devices = {TPU_REPLICATED_CORE_0 = ["/job:worker/replica:0/task:0/device:TPU:1", "/job:worker/replica:0/task:0/device:TPU:2"], TPU_REPLICATED_CORE_1 = ["/job:worker/replica:0/task:0/device:TPU:3", "/job:worker/replica:0/task:0/device:TPU:4"]}} {

      %2 = "tf.ReadVariableOp"(%arg15) : (tensor<*x!tf_type.resource<tensor<7x7x3x64xf32>>>) -> tensor<7x7x3x64xf32>