Fix performance issue with nanstd and nanvar

xtensor-stack · Apr 7, 2021 · 622647d · 622647d
1 parent 78aaac3
commit 622647d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/include/xtensor/xmath.hpp b/include/xtensor/xmath.hpp
@@ -2664,8 +2664,8 @@ namespace detail {
               XTL_REQUIRES(is_reducer_options<EVS>)>
     inline auto nanvar(E&& e, EVS es = EVS())
     {
-        decltype(auto) sc = detail::shared_forward<E>(e);
-        return nanmean<T>(square(sc - nanmean<T>(sc)), es);
+        auto cached_mean = nanmean<T>(e, es)();
+        return nanmean<T>(square(std::forward<E>(e) - std::move(cached_mean)), es);
     }
 
     template <class T = void, class E, class EVS = DEFAULT_STRATEGY_REDUCERS,
@@ -2702,7 +2702,8 @@ namespace detail {
         // note: forcing copy of first axes argument -- is there a better solution?
         auto axes_copy = axes;
         using result_type = typename std::conditional_t<std::is_same<T, void>::value, double, T>;
-        auto inner_mean = nanmean<result_type>(sc, std::move(axes_copy));
+        // always eval to prevent repeated evaluations in the next calls
+        auto inner_mean = eval(nanmean<result_type>(sc, std::move(axes_copy), evaluation_strategy::immediate));
 
         // fake keep_dims = 1
         auto keep_dim_shape = e.shape();