if (auto if_op = dyn_cast<TF::IfOp>(op)) {
      result = CanonicalizeFunctionalIfCase(
          op, {if_op.then_function(), if_op.else_function()}, if_op.getInput());
    } else if (auto case_op = dyn_cast<TF::CaseOp>(op)) {
      SmallVector<func::FuncOp, 4> branches;
      case_op.get_branch_functions(branches);
      result =
          CanonicalizeFunctionalIfCase(case_op, branches, case_op.getInput());