begin_if_op_iter = std::next(current_last_if_op_iter);
  }
}

// Checks whether a return index should be kept for `current_if_op` by checking
// for results in `if_op_segment`.
llvm::SmallVector<int, 4> GetReturnIndicesToKeep(
    TF::IfRegionOp current_if_op,
    const llvm::SmallVector<TF::IfRegionOp, 8>& if_op_segment) {
  llvm::SmallVector<int, 4> return_indices_to_keep;
  auto is_op_inside_IfRegions = [&](Operation* op) {