More generic implementation

xtensor-stack · Jan 1, 2024 · 995896b · 995896b
1 parent 1555d4d
commit 995896b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/include/xtensor/xbuilder.hpp b/include/xtensor/xbuilder.hpp
@@ -535,7 +535,7 @@ namespace xt
             using value_type = xtl::promote_type_t<typename std::decay_t<CT>::value_type...>;
 
             template <class It>
-            inline value_type access(const tuple_type& t, size_type axis, It first, It last) const
+            inline value_type access(const tuple_type& t, size_type axis, It first, It) const
             {
                 auto get_item = [&](auto& arr)
                 {
@@ -548,7 +548,8 @@ namespace xt
                         {
                             after_axis = true;
                         }
-                        const auto& stride = arr.strides()[i];
+                        const auto& shape = arr.shape();
+                        const size_t stride = std::accumulate(shape.begin() + i + 1, shape.end(), 1, std::multiplies<size_t>());
                         const auto len =  (*(first + i + after_axis));
                         offset += len * stride;
                     }

diff --git a/test/test_xbuilder.cpp b/test/test_xbuilder.cpp
@@ -466,9 +466,9 @@ namespace xt
         ASSERT_EQ(3, l(1, 1));
         ASSERT_EQ(3, l(2, 0));
 
-        // auto t = stack(xtuple(arange(3), arange(3, 6), arange(6, 9)));
-        // xarray<double> ar = {{0, 1, 2}, {3, 4, 5}, {6, 7, 8}};
-        // ASSERT_TRUE(t == ar);
+        auto t = stack(xtuple(arange(3), arange(3, 6), arange(6, 9)));
+        xarray<double> ar = {{0, 1, 2}, {3, 4, 5}, {6, 7, 8}};
+        ASSERT_TRUE(t == ar);
     }
 
     TEST(xbuilder, hstack)