// 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>):
    %1 = func.call @init(%arg4) : (tensor<4xf32>) -> tensor<7xf32>
    "tf.Yield"(%1) : (tensor<7xf32>) -> ()
  }, {

           while_op.getResults().take_front(while_region.getNumResults())))
    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));

  while_region.erase();
  return success();
}

// Transform GeneratorDatasetRegion to GeneratorDatasetOp.
LogicalResult RegionControlFlowToFunctional::ConvertGeneratorDatasetOp(
    SymbolTableCollection& symbol_table, GeneratorDatasetRegionOp regional) {
  mlir::MLIRContext* ctx = regional.getContext();

      build($_builder, $_state, resultTypes, operands, attributes);
    }]>];

  let hasCanonicalizer = 1;
}

def TF_GeneratorDatasetRegionOp : TF_Op<"GeneratorDatasetRegion",
      [AttrSizedOperandSegments,
       DeclareOpInterfaceMethods<RegionBranchOpInterface, [
           "areTypesCompatible",
           "getEntrySuccessorOperands",
           "getRegionInvocationBounds",