diff --git a/stan/math/prim/functor/reduce_sum_static.hpp b/stan/math/prim/functor/reduce_sum_static.hpp index 7ad8648aad8..09d570fa20e 100644 --- a/stan/math/prim/functor/reduce_sum_static.hpp +++ b/stan/math/prim/functor/reduce_sum_static.hpp @@ -50,9 +50,9 @@ inline auto reduce_sum_static(Vec&& vmapped, int grainsize, std::ostream* msgs, #ifdef STAN_THREADS return internal::reduce_sum_impl()(std::forward(vmapped), false, - grainsize, msgs, - std::forward(args)...); + ref_type_t...>()( + std::forward(vmapped), false, grainsize, msgs, + std::forward(args)...); #else if (vmapped.empty()) { return return_type(0); diff --git a/test/unit/math/prim/functor/reduce_sum_static_expression_test.cpp b/test/unit/math/prim/functor/reduce_sum_static_expression_test.cpp new file mode 100644 index 00000000000..e44bb64e8f1 --- /dev/null +++ b/test/unit/math/prim/functor/reduce_sum_static_expression_test.cpp @@ -0,0 +1,67 @@ +#include +#include + +#include +#include + +namespace stan { +namespace math { +namespace test { + +// Counts how many times the expression's coefficients are read. +template +struct counting_op { + std::atomic* count_; + + explicit counting_op(std::atomic* count) : count_(count) {} + + inline const T& operator()(const T& a) const { + ++(*count_); + return a; + } +}; + +struct expression_sum_lpdf { + template + inline auto operator()(const VecT& sub_slice, std::size_t start, + std::size_t end, std::ostream* msgs, + const EigExpr& shared) const { + // Read the expression once per call so the counter tracks calls, not terms. + stan::return_type_t shared_sum = stan::math::sum(shared); + stan::return_type_t sum = 0; + for (std::size_t i = 0; i < sub_slice.size(); ++i) { + sum += sub_slice[i] * shared_sum; + } + return sum; + } +}; + +} // namespace test +} // namespace math +} // namespace stan + +// https://github.com/stan-dev/math/issues/3304 +TEST(StanMathPrim_reduce_sum_static, eigen_expression_arg_evaluated_once) { + using stan::math::test::counting_op; + using stan::math::test::expression_sum_lpdf; + + Eigen::MatrixXd m = Eigen::MatrixXd::Ones(2, 2); + std::vector slice{1.0, 2.0, 3.0, 4.0}; + + std::atomic reduce_sum_count{0}; + counting_op op_dynamic(&reduce_sum_count); + double from_reduce_sum = stan::math::reduce_sum( + slice, 1, nullptr, m.block(0, 0, 1, 1).unaryExpr(op_dynamic)); + + std::atomic reduce_sum_static_count{0}; + counting_op op_static(&reduce_sum_static_count); + double from_reduce_sum_static + = stan::math::reduce_sum_static( + slice, 1, nullptr, m.block(0, 0, 1, 1).unaryExpr(op_static)); + + EXPECT_DOUBLE_EQ(from_reduce_sum, 10.0); + EXPECT_DOUBLE_EQ(from_reduce_sum_static, 10.0); + + EXPECT_EQ(1, reduce_sum_count.load()); + EXPECT_EQ(1, reduce_sum_static_count.load()); +} diff --git a/test/unit/math/rev/functor/reduce_sum_static_expression_test.cpp b/test/unit/math/rev/functor/reduce_sum_static_expression_test.cpp new file mode 100644 index 00000000000..317406cbc90 --- /dev/null +++ b/test/unit/math/rev/functor/reduce_sum_static_expression_test.cpp @@ -0,0 +1,45 @@ +#include +#include +#include + +#include + +namespace stan { +namespace math { +namespace test { + +struct expression_sum_lpdf { + template + inline auto operator()(const VecT& sub_slice, std::size_t start, + std::size_t end, std::ostream* msgs, + const EigExpr& shared) const { + stan::math::var sum = 0; + for (std::size_t i = 0; i < sub_slice.size(); ++i) { + sum += sub_slice[i] * stan::math::sum(shared); + } + return sum; + } +}; + +} // namespace test +} // namespace math +} // namespace stan + +// https://github.com/stan-dev/math/issues/3304 +TEST_F(AgradRev, StanMathRev_reduce_sum_static_eigen_expression_arg_gradient) { + using stan::math::var; + using stan::math::test::expression_sum_lpdf; + + Eigen::Matrix m = Eigen::MatrixXd::Ones(2, 2); + std::vector slice{1.0, 2.0, 3.0, 4.0}; + + var out = stan::math::reduce_sum_static( + slice, 1, nullptr, + m.block(0, 0, 1, 1).unaryExpr([](const var& a) { return a; })); + + EXPECT_DOUBLE_EQ(10.0, out.val()); + + out.grad(); + EXPECT_DOUBLE_EQ(10.0, m(0, 0).adj()); + EXPECT_DOUBLE_EQ(1.0, slice[0].adj()); +}