Skip to content

Commit d01be62

Browse files
committed
Improve performance of nanvar and nanmean
1 parent f3c11b2 commit d01be62

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

include/xtensor/xmath.hpp

+7-4
Original file line numberDiff line numberDiff line change
@@ -2671,8 +2671,8 @@ namespace detail {
26712671
XTL_REQUIRES(is_reducer_options<EVS>)>
26722672
inline auto nanvar(E&& e, EVS es = EVS())
26732673
{
2674-
decltype(auto) sc = detail::shared_forward<E>(e);
2675-
return nanmean<T>(square(sc - nanmean<T>(sc)), es);
2674+
auto cached_mean = nanmean<T>(e, es)();
2675+
return nanmean<T>(square(std::forward<E>(e) - std::move(cached_mean)), es);
26762676
}
26772677

26782678
template <class T = void, class E, class EVS = DEFAULT_STRATEGY_REDUCERS,
@@ -2709,7 +2709,8 @@ namespace detail {
27092709
// note: forcing copy of first axes argument -- is there a better solution?
27102710
auto axes_copy = axes;
27112711
using result_type = typename std::conditional_t<std::is_same<T, void>::value, double, T>;
2712-
auto inner_mean = nanmean<result_type>(sc, std::move(axes_copy));
2712+
// always eval to prevent repeated evaluations in the next calls
2713+
auto inner_mean = eval(nanmean<result_type>(sc, std::move(axes_copy)));
27132714

27142715
// fake keep_dims = 1
27152716
auto keep_dim_shape = e.shape();
@@ -2718,7 +2719,9 @@ namespace detail {
27182719
keep_dim_shape[el] = 1;
27192720
}
27202721
auto mrv = reshape_view<XTENSOR_DEFAULT_LAYOUT>(std::move(inner_mean), std::move(keep_dim_shape));
2721-
return nanmean<result_type>(square(cast<result_type>(sc) - std::move(mrv)), std::forward<X>(axes), es);
2722+
// note: otherwise the result is wrong with 'immediate' evaluation strategy
2723+
auto sc_shifted = eval(cast<result_type>(sc) - std::move(mrv));
2724+
return nanmean<result_type>(square(sc_shifted), std::forward<X>(axes), es);
27222725
}
27232726

27242727
/**

0 commit comments

Comments
 (0)