From 92700b3de5b55131b7c2232048b908d875e07031 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 1 Sep 2026 16:33:21 -0400 Subject: [PATCH 1/2] test(reduce_sum_static): add failing tests for missing ref_type_t wrapper Regression tests for #3304. Both fail against develop when built with STAN_THREADS; without threads reduce_sum_static bypasses reduce_sum_impl entirely, so the assertions hold trivially there. prim: passing an Eigen expression template as a shared argument to reduce_sum_static re-evaluates it once per TBB worker (4 reads with grainsize 1 over 4 terms) instead of once. reduce_sum, which already wraps Args with ref_type_t, reads it once. rev: the same call does not compile. The unwrapped CwiseUnaryOp reaches save_varis, which uses single-index coeff(i) on an expression without LinearAccessBit, tripping an Eigen static assertion. --- .../reduce_sum_static_expression_test.cpp | 67 +++++++++++++++++++ .../reduce_sum_static_expression_test.cpp | 45 +++++++++++++ 2 files changed, 112 insertions(+) create mode 100644 test/unit/math/prim/functor/reduce_sum_static_expression_test.cpp create mode 100644 test/unit/math/rev/functor/reduce_sum_static_expression_test.cpp 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()); +} From 435547c34376ed87229a04fd308ec9746538619c Mon Sep 17 00:00:00 2001 From: Trevin Chow Date: Sun, 19 Apr 2026 03:14:07 -0700 Subject: [PATCH 2/2] fix(reduce_sum_static): wrap Args with ref_type_t to materialize Eigen expressions reduce_sum has been wrapping shared Args with ref_type_t... since v4.8+ so that Eigen expression-template temporaries (e.g. M.row(0), matrix.col(i)) materialize before being handed off to TBB's parallel reducer. reduce_sum_static's STAN_THREADS branch was not updated in the same pass and still forwarded plain Args..., so callers passing a temporary Eigen expression as a shared argument could get a dangling reference once the temporary went out of scope. Mirror the reduce_sum wrapping exactly: internal::reduce_sum_impl<..., ref_type_t...>()(...) No other changes. ref_type_t is already in scope via stan/math/prim/meta.hpp (already included). Closes #3304 --- stan/math/prim/functor/reduce_sum_static.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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);