}

// 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,

  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