// CHECK: func.call @next
  // CHECK: ^
  // CHECK-SAME: tensor<6xf32>
  // CHECK-SAME: tensor<2xf32>
  // CHECK: func.call @finalize
  // CHECK-NOT: tf.GeneratorDataset
  %0 = "tf.GeneratorDataset"(%arg0, %arg1, %arg2, %arg3) {
      device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0",
      finalize_func = @finalize,
      init_func = @init,
      next_func = @next,

  tf_executor.graph {
    tf_executor.island {
      // CHECK: %{{.*}}, %[[CONTROL:[^ ,]*]] = tf_executor.island wraps "tf.GeneratorDataset"
      %gen0 = "tf.GeneratorDataset"(%str, %arg0, %arg1, %arg2, %arg3) {
        finalize_func = @__finalize_func_790,
        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,

  tf_executor.graph {
    tf_executor.island {
      // CHECK: %{{.*}}, %[[CONTROL:[^ ,]*]] = tf_executor.island wraps "tf.GeneratorDataset"
      %gen0 = "tf.GeneratorDataset"(%str, %arg0, %arg1, %arg2, %arg3) {
        finalize_func = @__finalize_func_790,
        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,

                                                                               \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("GeneratorDataset").Device(DEVICE).HostMemory("handle"),            \
      data::GeneratorDatasetOp);                                               \
  REGISTER_KERNEL_BUILDER(Name("PrefetchDataset")                              \

  tf_executor.graph {
    // expected-remark@above {{ID: 6}}
    %island = tf_executor.island {
        // expected-remark@above {{ID: 4}}
        // expected-remark@above {{Successors: {5}}}

}

// CHECK-LABEL: testGeneratorDatasetRegion
func.func @testGeneratorDatasetRegion(%arg0: tensor<4xf32>, %arg1: tensor<3xf32>, %arg2: tensor<!tf_type.resource>, %arg3: tensor<2xf32>) {
  // CHECK: "tf.GeneratorDataset"
  // CHECK-DAG: @init
  // CHECK-DAG: @next
  // CHECK-DAG: @finalize
  // CHECK: return
  %0 = "tf.GeneratorDatasetRegion"(%arg0, %arg1, %arg2, %arg3) ({
  ^bb0(%arg4: tensor<4xf32>):

           "getSuccessorRegions"
       ]>,
       SingleBlockImplicitTerminator<"YieldOp">,
       TF_GeneratorOpSideEffect,
      ]> {
  let summary = "Regional version of GeneratorDataset";

  let description = [{
Creates a dataset that invokes its 'next' region to generate elements. Conceptually,
within MLIR, we treat this op as if it fills a buffer with all the results right away,

                                %arg1: tensor<3xf32>,
                                %arg2: tensor<!tf_type.resource>,
                                %arg3: tensor<2xf32>) -> tensor<!tf_type.variant> {
  %0 = "tf.GeneratorDataset"(%arg0, %arg1, %arg2, %arg3) {
      device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0",
      finalize_func = @finalize,
      init_func = @init,
      next_func = @next,

  TF_DerivedOperandTypeAttr Tindices = TF_DerivedOperandTypeAttr<1>;
  TF_DerivedOperandTypeAttr Tparams = TF_DerivedOperandTypeAttr<0>;

  let hasVerifier = 1;
}

def TF_GeneratorDatasetOp : TF_Op<"GeneratorDataset", [AttrSizedOperandSegments, TF_GeneratorOpSideEffect]> {
  let summary = [{
Creates a dataset that invokes a function to generate elements.
  }];

  let arguments = (ins