%arg2: tensor<!tf_type.resource>,
                                %arg3: tensor<2xf32>) {
  // CHECK-NOT: tf.GeneratorDataset
  // CHECK: tf.GeneratorDatasetRegion
  // CHECK: ^
  // CHECK-SAME: tensor<4xf32>
  // CHECK: func.call @init
  // CHECK: ^
  // CHECK-SAME: tensor<7xf32>
  // CHECK-SAME: tensor<3xf32>
  // CHECK-NOT: tf_type.resource

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

func.func @testInvalidYieldOp(%arg0: f32) -> () {
  // expected-error @+1 {{'tf.Yield' op expects parent op to be one of 'tf.CaseRegion, tf.IfRegion, tf.WhileRegion, tf.GeneratorDatasetRegion'}}
  "tf.Yield"(%arg0) : (f32) -> ()
}

// -----

// Test valid tf.IfRegion operation
// CHECK-LABEL: func @testValidIfRegionOp