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.

Author: gouarin

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,20 +1797,23 @@ 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 <class F, class CT, class X, class O>
-    inline auto xreducer_stepper<F, CT, X, O>::aggregate_impl(size_type dim, std::false_type) const -> reference
+    inline auto xreducer_stepper<F, CT, X, O>::aggregate_impl(size_type dim) const -> reference
     {
         // reference can be std::array, hence the {} initializer
         reference res = {};
         size_type index = axis(dim);
         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 <class F, class CT, class X, class O>
-    inline auto xreducer_stepper<F, CT, X, O>::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<reference>(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 <class F, class CT, class X, class O>
     inline auto xreducer_stepper<F, CT, X, O>::get_substepper_begin() const -> substepper_type
     {

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<int> a = xt::reshape_view(xt::arange<int>(2 * 3 * 4 * 5), {2, 3, 4, 5});
+        xt::xarray<int, xt::layout_type::column_major> c = a;
+
+        // Every non-empty set of axes, including non-adjacent ones.
+        for (std::size_t mask = 1; mask < 16; ++mask)
+        {
+            std::vector<std::size_t> 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<int> lazy = xt::sum(a, axes, xt::keep_dims);
+            EXPECT_EQ(lazy, expected);
+
+            xt::xarray<int> lazy_init = xt::sum(a, axes, xt::keep_dims | initial(7));
+            EXPECT_EQ(lazy_init, expected + 7);
+
+            xt::xarray<int> lazy_col = xt::sum(c, axes, xt::keep_dims);
+            EXPECT_EQ(lazy_col, expected);
+        }
+    }
+
     TEST(xreducer, ones_first)
     {
         auto a = xt::ones<int>(std::vector<int>({1, 1, 2, 4}));