}
    tf_executor.fetch
  }
  func.return
}
// CHECK: "tf.opA"
// device = "/TPU:0"
// CHECK: _parallel_execution_ids = "r0:0"
// CHECK: "tf.opA"
// device = "/TPU:0"
// CHECK: _parallel_execution_ids = "r0:1"
// CHECK: "tf.opA"
// device = "/TPU:1"
// CHECK: _parallel_execution_ids = "r1:0"
// CHECK: "tf.opA"
// device = "/TPU:1"

func.func @reorder_single_op_program_key() {
  // CHECK: %[[COMPILE_LAUNCH:[0-9]*]]:2 = "tf_device.launch"
  // CHECK: TPUCompileMlir
  // CHECK: "tf.OpA"(%[[COMPILE_LAUNCH]]#1
  "tf_device.launch"() ({
    %cst_0 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
    "tf.OpA"(%cst_0) { mini_batch_splits = ""} : (tensor<1x!tf_type.string>) -> ()
    tf_device.return

  tf_executor.graph {
    %1:2 = tf_executor.island {
      %3 = "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
      %4 = "tf.opB"(%3) : (tensor<i1>) -> tensor<i1>
      tf_executor.yield %3 : tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK-NEXT: %[[OP_A:[0-9]*]] = "tf.opA"(%[[ARG_0]])
// CHECK-NEXT: "tf.opB"(%[[OP_A]])
// CHECK-NEXT: return

    %1:2 = tf_executor.island wraps "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
    %2:2 = tf_executor.island(%1#1) wraps "tf.opB"(%arg1) : (tensor<i1>) -> tensor<i1>
    tf_executor.fetch %1#0, %2#0 : tensor<i1>, tensor<i1>
  }
  func.return %0#0, %0#1 : tensor<i1>, tensor<i1>
}

// CHECK:        %[[ISLAND:.*]]:2, %{{.*}} = tf_executor.island {
// CHECK-NEXT:     %[[OP_A:[0-9]*]] = "tf.opA"(%[[ARG_0]])

    %2 = "tf.opA"(%1, %ri) : (tensor<?xi32>, tensor<*xf32>) -> tensor<*xi32>
    tf_device.return %1, %2 : tensor<?xi32>, tensor<*xi32>
  }
  func.return
}

// CHECK: %[[SHAPE:[0-9]*]] = "tf.Shape"(%[[ARG_0]])
// CHECK: tf_device.replicate([%[[ARG_0]], %[[ARG_1]]] as %[[RI:[a-z0-9]*]]: tensor<*xf32>)
// CHECK:   %[[OP_A:[0-9]*]] = "tf.opA"(%[[SHAPE]], %[[RI]])

    %out, %c1 = tf_executor.island wraps "tf.opA"(%ri) {_tpu_replicate="cluster"}: (tensor<i32>) -> tensor<i32>
    // expected-warning @+1 {{TF2XLA TPU bridge input check: cluster op = tf.opA with cluster = cluster has successor as non cluster op tf.opB}}
    %out2, %c2 = tf_executor.island wraps "tf.opB"(%out) : (tensor<i32>) -> tensor<i32>

// CHECK:       [[RES_1:%.*]] = "tf.Add"([[CST_1]], [[CST_0]])
// CHECK:       [[RES_2:%.*]] = "tf.Mul"([[RES_1]], [[CST_1]])
// CHECK:       tf.WhileRegion
// CHECK:       ^bb0
// CHECK:       tf.OpA
// CHECK:       ^bb0([[ARG_2:%[a-zA-Z0-9_]+]]
// CHECK-SAME:  [[ARG_3:%[a-zA-Z0-9_]+]]: tensor<i32>)
// CHECK-NEXT:  [[RES_3:%.*]] = "tf.AddV2"([[ARG_2]], [[RES_1]])
// CHECK-NEXT:  [[RES_4:%.*]] = "tf.Div"([[ARG_3]], [[RES_2]])

    For example, the following:

    ```mlir
    func @single_op_launch() {
      %a = "tf.opA"() {device = "CPU:0"} : () -> tensor<i1>
      return %a
    }
    ```

    will be transformed into:

    ```mlir
    func @single_op_launch() {
      %1 = tf_device.launch() ( {
        %a = "tf.opA"() : () -> tensor<i1>
        tf_device.return %a
      }) {device = "CPU:0"} : () -> tensor<i1>

  //
  // Two ops with the same groups and different branches are considered
  // parallel so are not made dependent. For example if `OpA` has parallel_ids
  //   `{group0:branch0, group1:branch0}`
  // and `OpB` has parallel_ids
  //   `{group0:branch1, graph1:branch0}`
  // then `OpA` and `OpB` are executed in parallel because `group0` is common
  // with a different branch.
  //

  %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  // CHECK:      [[COMPUTATION:%.+]] = "tf_device.parallel_execute"()
  // CHECK:      "tf.OpA"([[READ0]])
  %2 = "tf_device.parallel_execute"() ({
    "tf_device.launch"() ({
      "tf.OpA"(%1) : (tensor<i32>) -> ()
      tf_device.return
    }) {device = "TPU_REPLICATED_HOST_0"} : () -> ()
    tf_device.return
  }, {