%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
  }

// This pass forms `tf_executor.island` per region of
// `tf_device.parallel_execute`.
//
// For example, the following:
//
//  %0 = tf_executor.island {
//    tf_executor.yield
//  }
//  %1:2 = tf_executor.island {
//    %2 = "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
//      tf_executor.yield %2 : tensor<i1>
//  }
//  %3:2 = tf_executor.island(%0) {
//    %4 = "tf_device.parallel_execute"() ({

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>

    dependencies.

    Operations in this dialect return a value of type `!tf_executor.control` as
    last returned value (exceptions are `tf_executor.NextIteration.graph`,
    `tf_executor.NextIteration.sink` and `tf_executor.fetch` which don’t return any
    value).
}];

  let cppNamespace = "::mlir::tf_executor";
}

// Control type.
def TfeControlType : Type<CPred<"$_self.isa<ControlType>()">, "control">,

          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"

    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<"value"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
      %outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<"value"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>