reduce_sum.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_FWD_FUNCTOR_REDUCE_SUM_HPP
#define STAN_MATH_FWD_FUNCTOR_REDUCE_SUM_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
 
#include <algorithm>
#include <tuple>
#include <vector>
 
namespace stan {
namespace math {
 
namespace internal {
template <typename ReduceFunction, typename Enable, typename ReturnType,
          typename Vec, typename... Args>
struct reduce_sum_impl;
 
template <typename ReduceFunction, typename Enable, typename ReturnType,
          typename Vec, typename... Args>
struct reduce_sum_impl {
  inline return_type_t<Vec, Args...> operator()(Vec&& vmapped,
                                                bool auto_partitioning,
                                                int grainsize,
                                                std::ostream* msgs,
                                                Args&&... args) const {
    if (vmapped.empty()) {
      return 0.0;
    }
 
    if (auto_partitioning) {
      return ReduceFunction()(std::forward<Vec>(vmapped), 0, vmapped.size() - 1,
                              msgs, std::forward<Args>(args)...);
    } else {
      return_type_t<Vec, Args...> sum = 0.0;
      for (size_t i = 0; i < (vmapped.size() + grainsize - 1) / grainsize;
           ++i) {
        size_t start = i * grainsize;
        size_t end = std::min((i + 1) * grainsize, vmapped.size()) - 1;
 
        std::decay_t<Vec> sub_slice;
        sub_slice.reserve(end - start + 1);
        for (size_t i = start; i <= end; ++i) {
          sub_slice.emplace_back(vmapped[i]);
        }
 
        sum += ReduceFunction()(std::forward<Vec>(sub_slice), start, end, msgs,
                                std::forward<Args>(args)...);
      }
      return sum;
    }
  }
};
 
}  // namespace internal
 
}  // namespace math
}  // namespace stan
 
#endif