Skip to content

Use one aggregate_impl for reducers with and without keep_dims #2959

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gouarin wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Use one aggregate_impl for reducers with and without keep_dims #2959

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 8 additions & 46 deletions include/xtensor/reducers/xreducer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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
Expand All @@ -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
{
Expand Down
29 changes: 29 additions & 0 deletions test/test_xreducer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}));
Expand Down
Loading