From 30bc8ea3c26aeb4fa5c91e034edf6e0922d80d47 Mon Sep 17 00:00:00 2001 From: Loic Gouarin Date: Thu, 29 Oct 2026 18:50:04 +0100 Subject: [PATCH] Use one aggregate_impl for reducers with and without keep_dims xreducer_stepper::aggregate_impl had two overloads. The one used without keep_dims walks the reduced axes, m_axes[0] to m_axes.back(). The keep_dims one walked every dimension of the expression, looked each one up in m_axes with std::find, and recursed through the dimensions that are not reduced without doing anything there. Both do the same steps on the underlying stepper: keep_dims only changes the shape of the result, which the stepper does not use here. The extra size != 0 test of the keep_dims overload is also redundant, since aggregate() returns the initial value for empty expressions before calling aggregate_impl. The keep_dims overload is removed. GCC 13 at -O3 inlined that recursion and warned about m_e.shape()[dim] for dim == dimension() (-Warray-bounds), a path that the axis validation in the constructor rules out. The warning goes away with the overload. The new test compares lazy keep_dims sums with immediate ones for every set of axes of a 4-D array, in row and column major order and with an initial value. --- include/xtensor/reducers/xreducer.hpp | 54 ++++----------------------- test/test_xreducer.cpp | 29 ++++++++++++++ 2 files changed, 37 insertions(+), 46 deletions(-) diff --git a/include/xtensor/reducers/xreducer.hpp b/include/xtensor/reducers/xreducer.hpp index e99522b5e..e491082fe 100644 --- a/include/xtensor/reducers/xreducer.hpp +++ b/include/xtensor/reducers/xreducer.hpp @@ -1123,8 +1123,7 @@ namespace xt reference initial_value() const; reference aggregate(size_type dim) const; - reference aggregate_impl(size_type dim, /*keep_dims=*/std::false_type) const; - reference aggregate_impl(size_type dim, /*keep_dims=*/std::true_type) const; + reference aggregate_impl(size_type dim) const; substepper_type get_substepper_begin() const; size_type get_dim(size_type dim) const noexcept; @@ -1789,7 +1788,7 @@ namespace xt } else { - res = aggregate_impl(dim, typename O::keep_dims()); + res = aggregate_impl(dim); if (O::has_initial_value && dim == 0) { res = m_reducer->m_merge(m_reducer->m_options.initial_value, res); @@ -1798,8 +1797,11 @@ namespace xt return res; } + // Reduces along the axes m_axes[dim], ..., m_axes.back() at the current position + // of the stepper, then moves the stepper back along m_axes[dim]. Dimensions that + // are not reduced are not visited, with or without keep_dims. template - inline auto xreducer_stepper::aggregate_impl(size_type dim, std::false_type) const -> reference + inline auto xreducer_stepper::aggregate_impl(size_type dim) const -> reference { // reference can be std::array, hence the {} initializer reference res = {}; @@ -1807,11 +1809,11 @@ namespace xt size_type size = shape(index); if (dim != m_reducer->m_axes.size() - 1) { - res = aggregate_impl(dim + 1, typename O::keep_dims()); + res = aggregate_impl(dim + 1); for (size_type i = 1; i != size; ++i) { m_stepper.step(index); - res = m_reducer->m_merge(res, aggregate_impl(dim + 1, typename O::keep_dims())); + res = m_reducer->m_merge(res, aggregate_impl(dim + 1)); } } else @@ -1827,46 +1829,6 @@ namespace xt return res; } - template - inline auto xreducer_stepper::aggregate_impl(size_type dim, std::true_type) const -> reference - { - // reference can be std::array, hence the {} initializer - reference res = {}; - auto ax_it = std::find(m_reducer->m_axes.begin(), m_reducer->m_axes.end(), dim); - if (ax_it != m_reducer->m_axes.end()) - { - size_type index = dim; - size_type size = m_reducer->m_e.shape()[index]; - if (ax_it != m_reducer->m_axes.end() - 1 && size != 0) - { - res = aggregate_impl(dim + 1, typename O::keep_dims()); - for (size_type i = 1; i != size; ++i) - { - m_stepper.step(index); - res = m_reducer->m_merge(res, aggregate_impl(dim + 1, typename O::keep_dims())); - } - } - else - { - res = m_reducer->m_reduce(static_cast(m_reducer->m_init()), *m_stepper); - for (size_type i = 1; i != size; ++i) - { - m_stepper.step(index); - res = m_reducer->m_reduce(res, *m_stepper); - } - } - m_stepper.reset(index); - } - else - { - if (dim < m_reducer->m_e.dimension()) - { - res = aggregate_impl(dim + 1, typename O::keep_dims()); - } - } - return res; - } - template inline auto xreducer_stepper::get_substepper_begin() const -> substepper_type { diff --git a/test/test_xreducer.cpp b/test/test_xreducer.cpp index aa0a824b5..b3107235f 100644 --- a/test/test_xreducer.cpp +++ b/test/test_xreducer.cpp @@ -1017,6 +1017,35 @@ namespace xt EXPECT_TRUE(xt::all(equal(rex0, 0.))); } + TEST(xreducer, keep_dims_lazy_all_axes) + { + xt::xarray a = xt::reshape_view(xt::arange(2 * 3 * 4 * 5), {2, 3, 4, 5}); + xt::xarray c = a; + + // Every non-empty set of axes, including non-adjacent ones. + for (std::size_t mask = 1; mask < 16; ++mask) + { + std::vector axes; + for (std::size_t d = 0; d < 4; ++d) + { + if (mask & (std::size_t(1) << d)) + { + axes.push_back(d); + } + } + auto expected = xt::sum(a, axes, xt::keep_dims | xt::evaluation_strategy::immediate); + + xt::xarray lazy = xt::sum(a, axes, xt::keep_dims); + EXPECT_EQ(lazy, expected); + + xt::xarray lazy_init = xt::sum(a, axes, xt::keep_dims | initial(7)); + EXPECT_EQ(lazy_init, expected + 7); + + xt::xarray lazy_col = xt::sum(c, axes, xt::keep_dims); + EXPECT_EQ(lazy_col, expected); + } + } + TEST(xreducer, ones_first) { auto a = xt::ones(std::vector({1, 1, 2, 4}));