Bugfix: calling argmax(a, 0) with a.dimension() == 1 resulted in argmin(a, 0)

tdegeus · tdegeus · commit fc39f69962a0 · 2022-07-26T14:22:04.000+02:00
diff --git a/include/xtensor/xsort.hpp b/include/xtensor/xsort.hpp
@@ -777,14 +777,6 @@ namespace xt
             using result_type = typename argfunc_result_type<E>::type;
             using result_shape_type = typename result_type::shape_type;
 
-            if (e.dimension() == 1)
-            {
-                auto begin = e.template begin<L>();
-                auto end = e.template end<L>();
-                std::size_t i = static_cast<std::size_t>(std::distance(begin, std::min_element(begin, end)));
-                return xtensor<size_t, 0>{i};
-            }
-
             result_shape_type alt_shape;
             xt::resize_container(alt_shape, e.dimension() - 1);
 
@@ -854,6 +846,11 @@ namespace xt
     template <layout_type L = XTENSOR_DEFAULT_TRAVERSAL, class E>
     inline auto argmin(const xexpression<E>& e, std::ptrdiff_t axis)
     {
+        if (e.dimension() == 1)
+        {
+            return argmin(e);
+        }
+
         using value_type = typename E::value_type;
         auto&& ed = eval(e.derived_cast());
         std::size_t ax = normalize_axis(ed.dimension(), axis);
@@ -884,6 +881,11 @@ namespace xt
     template <layout_type L = XTENSOR_DEFAULT_TRAVERSAL, class E>
     inline auto argmax(const xexpression<E>& e, std::ptrdiff_t axis)
     {
+        if (e.dimension() == 1)
+        {
+            return argmax(e);
+        }
+
         using value_type = typename E::value_type;
         auto&& ed = eval(e.derived_cast());
         std::size_t ax = normalize_axis(ed.dimension(), axis);
diff --git a/test/test_xsort.cpp b/test/test_xsort.cpp
@@ -190,6 +190,10 @@ namespace xt
         EXPECT_EQ(ex, argmin(xa));
         EXPECT_EQ(ex_2, argmin(xa, 0));
         EXPECT_EQ(ex_3, argmin(xa, 1));
+
+        xtensor<double, 1> ya = {0, 1, 2, 3, 4};
+        EXPECT_EQ(0, argmin(ya)());
+        EXPECT_EQ(0, argmin(ya, 0)());
     }
 
     TEST(xsort, argmax)
@@ -218,6 +222,10 @@ namespace xt
 
         xtensor<std::size_t, 2> ex_6 = {{0, 0, 0, 0}, {0, 0, 0, 0}};
         EXPECT_EQ(ex_6, argmax(c, 1));
+
+        xtensor<double, 1> ya = {0, 1, 2, 3, 4};
+        EXPECT_EQ(4, argmax(ya)());
+        EXPECT_EQ(4, argmax(ya, 0)());
     }
 
     TEST(xsort, sort_large_prob)
