llvm::SmallVector<llvm::SmallVector<int, 4>> GetReturnIndicesVec(
    const llvm::SmallVector<TF::IfRegionOp, 8>& if_op_segment) {
  llvm::SmallVector<llvm::SmallVector<int, 4>> return_indices_vec;
  for (auto it = if_op_segment.begin(); it != if_op_segment.end(); ++it) {
    llvm::SmallVector<int, 4> indices_to_keep_vec =
        GetReturnIndicesToKeep(*it, if_op_segment);
    return_indices_vec.push_back(indices_to_keep_vec);
  }
  return return_indices_vec;