auto generator_region =
      OpBuilder(generator_op)
          .create<TF::GeneratorDatasetRegionOp>(
              generator_op.getLoc(), generator_op->getResultTypes(),
              generator_op.getInitFuncOtherArgs(),
              generator_op.getNextFuncOtherArgs(),
              generator_op.getFinalizeFuncOtherArgs(),
              generator_op.getOutputTypes(), generator_op.getOutputShapes(),

                               /*allow_return_of_existing=*/true);

  auto new_op = OpBuilder(regional).create<TF::GeneratorDatasetOp>(
      regional.getLoc(), regional->getResultTypes(),
      regional.getInitFuncOtherArgs(), regional.getNextFuncOtherArgs(),
      regional.getFinalizeFuncOtherArgs(), SymbolRefAttr::get(ctx, init_name),
      SymbolRefAttr::get(ctx, next_name),

  if (point.isParent()) {
    // The op itself doesn't branch back to itself.
    return ::mlir::OperandRange(end, end);
  } else if (point.getRegionOrNull() == &getInit()) {
    return getInitFuncOtherArgs();
  } else if (point.getRegionOrNull() == &getNext()) {
    return getNextFuncOtherArgs();
  } else /* finalize region */ {
    return getFinalizeFuncOtherArgs();
  }
}