%c = "tf.opC"() : () -> tensor<i1>
      tf_executor.yield %a, %launch#0, %launch#1, %c : tensor<i1>, tensor<f32>, tensor<i32>, tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK-NOT: tf_executor.island
// CHECK:      %[[A:.*]], {{.*}} = tf_executor.island wraps "tf.opA"
// CHECK:      %[[B:.*]]:2, {{.*}} = tf_executor.island wraps "tf.opB"(%[[A]])
// CHECK-SAME: device = "CPU:0"

// CHECK:     i: 2
// CHECK:   }
// CHECK: }

func.func @main() {
  tf_executor.graph {
    %dim:2 = tf_executor.island wraps "tf.Const"() {dtype = "tftype$DT_INT32", value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %input:2 = tf_executor.island wraps "tf.Const"() {dtype = "tftype$DT_INT32", value = dense<1.0> : tensor<4x6xf32>} : () -> tensor<4x6xf32>

func.func @colocate_split_with_pred() {
  // CHECK: Split
  // CHECK-SAME: _class = ["loc:@class"]
  tf_executor.graph {
    %c, %control0 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %a, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>)

    %2:2 = tf_executor.Enter %1#0 frame "while/while_context" parallel_iterations 10 : (tensor<!tf_type.int32ref>) -> (tensor<*x!tf_type.int32ref>, !tf_executor.control) {device = "", T = "tfdtype$DT_INT32"} loc("while/Enter")
    %3:3 = tf_executor.Merge %2#0, %0#0 : tensor<*x!tf_type.int32ref> {device = "", N = 2, T = "tfdtype$DT_INT32"} loc("while/Merge")

  tf_executor.graph {
    // CHECK: tf_executor.island
    // CHECK-NOT: wraps
    %0:3 = tf_executor.island {
      %1:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>) loc("identity@some_function")
      tf_executor.yield %1#0, %1#1 : tensor<i32>, tensor<f32> loc(unknown)
    } loc("identity@some_function")
    tf_executor.fetch
  }
  func.return
}

// NextIteration.Source and NextIteration.Sink.
// CHECK-LABEL: func @main
// CHECK: tf_executor.NextIteration.Source
// CHECK: tf_executor.NextIteration.Source
// CHECK: tf_executor.NextIteration.Sink
// CHECK: tf_executor.NextIteration.Sink
func.func @main() -> (tensor<1xi32>, tensor<1xi32>) {
  %0, %1 = tf_executor.graph {
    %output_1, %token_1, %control_1 = tf_executor.NextIteration.Source : tensor<1xi32> {T = i32, device = ""}

    tf_executor.graph {
      %control = tf_executor.island {
        "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @partitioned}> : (tensor<*x!tf_type.resource>) -> ()
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }

func.func @main() -> tensor<2x3xi32> {
  %graph = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<5> : tensor<2x3xi32>} : () -> tensor<2x3xi32> loc("Const0")
    %1:2 = tf_executor.island wraps "tf.Const"() {value = dense<4.2> : tensor<4x5xf32>} : () -> tensor<4x5xf32> loc("Const1")

    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<32xf32>} : () -> tensor<32xf32>
      %outputs_0, %control_1 = tf_executor.island(%control) wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<5x5x3x32xf32>} : () -> tensor<5x5x3x32xf32>

          tf_device.return
        }) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : () -> ()
        tf_executor.yield
      }
      tf_executor.fetch
    }
    func.return
  }

  // CHECK:      tf_executor.island
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
  // CHECK-SAME:   device_ordinal = 0
  // CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"