def TF_IfRegionOp : TF_Op<"IfRegion",
      [SingleBlockImplicitTerminator<"YieldOp">, NoRegionArguments,
       DeclareOpInterfaceMethods<RegionBranchOpInterface, [
           "areTypesCompatible",
           "getEntrySuccessorOperands",
           "getRegionInvocationBounds",
           "getSuccessorRegions"
       ]>
      ]> {
  let summary = "output = cond ? then_branch output : else_branch output";

  let description = [{

  invocationBounds.emplace_back(InvocationBounds::getUnknown());  // next
  invocationBounds.emplace_back(InvocationBounds(1, 1));          // finalize
}

OperandRange GeneratorDatasetRegionOp::getEntrySuccessorOperands(
    RegionBranchPoint point) {
  auto end = this->getOperation()->operand_end();
  if (point.isParent()) {
    // The op itself doesn't branch back to itself.
    return ::mlir::OperandRange(end, end);

// WhileRegionOp RegionBranchOpInterface
//===----------------------------------------------------------------------===//

OperandRange WhileRegionOp::getEntrySuccessorOperands(
    RegionBranchPoint point) {
  if (point.isParent()) {
    // WhileRegionOp branches to the condition, which branches to the body. But