feat: add multi-axis support for xt::roll with optimized pointer arithmetic

Author: f14XuanLv

include/xtensor/misc/xmanipulation.hpp

@@ -1033,6 +1033,77 @@ namespace xt
             }
             return to;
         }
+
+        /**
+         * Multi-axis roll using pointer arithmetic.
+         * Algorithm extended from single-axis roll above.
+         *
+         * @param to output iterator
+         * @param from input iterator
+         * @param shifts normalized shift values for each axis (size == shape.size())
+         * @param shape the shape of the tensor
+         * @param M current dimension being processed
+         */
+        template <class To, class From, class Shifts, class S>
+        To roll_multi(To to, From from, const Shifts& shifts, const S& shape, std::size_t M)
+        {
+            std::ptrdiff_t dim = static_cast<std::ptrdiff_t>(shape[M]);
+            std::ptrdiff_t offset = std::accumulate(
+                shape.begin() + M + 1,
+                shape.end(),
+                std::ptrdiff_t(1),
+                std::multiplies<std::ptrdiff_t>()
+            );
+            std::ptrdiff_t shift = shifts[M];
+
+            if (shape.size() == M + 1)
+            {
+                // Innermost dimension: direct copy
+                if (shift != 0)
+                {
+                    const auto split = from + (dim - shift) * offset;
+                    for (auto iter = split, end = from + dim * offset; iter != end; iter += offset, ++to)
+                    {
+                        *to = *iter;
+                    }
+                    for (auto iter = from, end = split; iter != end; iter += offset, ++to)
+                    {
+                        *to = *iter;
+                    }
+                }
+                else
+                {
+                    for (auto iter = from, end = from + dim * offset; iter != end; iter += offset, ++to)
+                    {
+                        *to = *iter;
+                    }
+                }
+            }
+            else
+            {
+                // Recursive case: process current dimension, then recurse
+                if (shift != 0)
+                {
+                    const auto split = from + (dim - shift) * offset;
+                    for (auto iter = split, end = from + dim * offset; iter != end; iter += offset)
+                    {
+                        to = roll_multi(to, iter, shifts, shape, M + 1);
+                    }
+                    for (auto iter = from, end = split; iter != end; iter += offset)
+                    {
+                        to = roll_multi(to, iter, shifts, shape, M + 1);
+                    }
+                }
+                else
+                {
+                    for (auto iter = from, end = from + dim * offset; iter != end; iter += offset)
+                    {
+                        to = roll_multi(to, iter, shifts, shape, M + 1);
+                    }
+                }
+            }
+            return to;
+        }
     }
 
     /**
@@ -1073,6 +1144,87 @@ namespace xt
         return cpy;
     }
 
+    /**
+     * Roll an expression along multiple axes.
+     *
+     * Elements that roll beyond the last position are re-introduced at the first.
+     * This function does not change the input expression.
+     *
+     * @ingroup xt_xmanipulation
+     * @param e the input xexpression
+     * @param shifts container of shift values for each axis
+     * @param axes container of axes along which elements are shifted
+     * @return a rolled copy of the input expression
+     *
+     * @note shifts and axes must have the same size. Each element in shifts
+     *       corresponds to the shift amount for the axis at the same position in axes.
+     * @note If the same axis appears multiple times, the shifts are accumulated.
+     * @note Negative axis indices are supported (e.g., -1 refers to the last axis).
+     *
+     * Example:
+     * @code
+     * xt::xarray<int> a = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
+     * auto result = xt::roll(a, {1, 2}, {0, 1});  // roll 1 on axis 0, roll 2 on axis 1
+     * @endcode
+     */
+    template <class E, class S, class X, XTL_REQUIRES(std::negation<xtl::is_integral<std::decay_t<S>>>)>
+    inline auto roll(E&& e, S&& shifts, X&& axes)
+    {
+        XTENSOR_ASSERT(std::size(shifts) == std::size(axes));
+
+        if (std::size(shifts) == 0)
+        {
+            return empty_like(e) = e;
+        }
+
+        const auto dim = e.dimension();
+        auto cpy = empty_like(e);
+        const auto& shape = cpy.shape();
+
+        // Accumulate shifts per axis (like NumPy)
+        // and normalize to positive values in range [0, axis_size)
+        std::vector<std::ptrdiff_t> total_shifts(dim, 0);
+        auto shift_it = std::begin(shifts);
+        auto axis_it = std::begin(axes);
+        for (; shift_it != std::end(shifts); ++shift_it, ++axis_it)
+        {
+            auto ax = normalize_axis(dim, static_cast<std::ptrdiff_t>(*axis_it));
+            total_shifts[ax] += static_cast<std::ptrdiff_t>(*shift_it);
+        }
+
+        // Normalize shifts to positive values
+        for (std::size_t ax = 0; ax < dim; ++ax)
+        {
+            auto axis_size = static_cast<std::ptrdiff_t>(shape[ax]);
+            if (axis_size > 0)
+            {
+                total_shifts[ax] = ((total_shifts[ax] % axis_size) + axis_size) % axis_size;
+            }
+        }
+
+        detail::roll_multi(cpy.begin(), e.begin(), total_shifts, shape, 0);
+        return cpy;
+    }
+
+    /**
+     * Roll an expression along multiple axes (C-style array overload).
+     *
+     * @ingroup xt_xmanipulation
+     * @param e the input xexpression
+     * @param shifts C-style array of shift values
+     * @param axes C-style array of axes
+     * @return a roll of the input expression
+     */
+    template <class E, class S, std::size_t N, class I, std::size_t M>
+    inline auto roll(E&& e, const S (&shifts)[N], const I (&axes)[M])
+    {
+        return roll(
+            std::forward<E>(e),
+            xtl::forward_sequence<std::array<S, N>, decltype(shifts)>(shifts),
+            xtl::forward_sequence<std::array<I, M>, decltype(axes)>(axes)
+        );
+    }
+
     /****************************
      * repeat implementation    *
      ****************************/

test/test_xmanipulation.cpp

@@ -504,6 +504,80 @@ namespace xt
         ASSERT_EQ(expected8, xt::roll(e2, -2, /*axis*/ 2));
     }
 
+    TEST(xmanipulation, roll_multi_axis)
+    {
+        // Test 1: Basic 2D multi-axis roll
+        xarray<double> e1 = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
+
+        // Roll 1 on axis 0, 2 on axis 1
+        xarray<double> expected1 = {{8, 9, 7}, {2, 3, 1}, {5, 6, 4}};
+        ASSERT_EQ(expected1, xt::roll(e1, {1, 2}, {0, 1}));
+
+        // Verify equivalence with sequential single-axis rolls
+        auto sequential_result = xt::roll(xt::roll(e1, 1, 0), 2, 1);
+        ASSERT_EQ(expected1, sequential_result);
+
+        // Test 2: std::array as input
+        std::array<int, 2> shifts = {1, 2};
+        std::array<int, 2> axes = {0, 1};
+        ASSERT_EQ(expected1, xt::roll(e1, shifts, axes));
+
+        // Test 3: std::vector as input
+        std::vector<int> shifts_v = {1, 2};
+        std::vector<int> axes_v = {0, 1};
+        ASSERT_EQ(expected1, xt::roll(e1, shifts_v, axes_v));
+
+        // Test 4: Negative shifts
+        xarray<double> expected4 = {{6, 4, 5}, {9, 7, 8}, {3, 1, 2}};
+        ASSERT_EQ(expected4, xt::roll(e1, {-1, -2}, {0, 1}));
+
+        // Test 5: Negative axis indices
+        ASSERT_EQ(expected1, xt::roll(e1, {1, 2}, {-2, -1}));
+
+        // Test 6: 3D array
+        xarray<double> e3 = {{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}};
+        xarray<double> expected6 = {{{8, 7}, {6, 5}}, {{4, 3}, {2, 1}}};
+        ASSERT_EQ(expected6, xt::roll(e3, {1, 1, 1}, {0, 1, 2}));
+
+        // Test 7: Single axis via multi-axis interface (should match single-axis version)
+        xarray<double> expected7 = {{7, 8, 9}, {1, 2, 3}, {4, 5, 6}};
+        ASSERT_EQ(expected7, xt::roll(e1, {1}, {0}));
+        ASSERT_EQ(expected7, xt::roll(e1, 1, 0));
+
+        // Test 8: Empty shifts (should return copy of original)
+        std::vector<int> empty_shifts;
+        std::vector<int> empty_axes;
+        ASSERT_EQ(e1, xt::roll(e1, empty_shifts, empty_axes));
+
+        // Test 9: Large shift values (should wrap around)
+        // shift of 10 on axis 0 (size 3) is equivalent to shift of 1
+        ASSERT_EQ(xt::roll(e1, {1}, {0}), xt::roll(e1, {10}, {0}));
+
+        // Test 10: Same axis appears multiple times (shifts accumulate)
+        // NumPy: np.roll(a, (1, 2), axis=(0, 0)) equals np.roll(a, 3, axis=0)
+        xarray<double> expected10 = xt::roll(e1, 3, 0);
+        ASSERT_EQ(expected10, xt::roll(e1, {1, 2}, {0, 0}));
+
+        // Test 11: xtensor (fixed dimension) instead of xarray
+        xt::xtensor<double, 2> t1 = {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}};
+        xt::xtensor<double, 2> t_expected = {{8, 9, 7}, {2, 3, 1}, {5, 6, 4}};
+        ASSERT_EQ(t_expected, xt::roll(t1, {1, 2}, {0, 1}));
+
+        // Test 12: 1D array multi-axis operation (only one axis)
+        xarray<double> e1d = {1, 2, 3, 4, 5};
+        xarray<double> expected1d = {4, 5, 1, 2, 3};
+        ASSERT_EQ(expected1d, xt::roll(e1d, {2}, {0}));
+
+        // Test 13: Column-major layout (result should be layout-independent)
+        xarray<double, layout_type::column_major> cm = {{1, 2, 3}, {4, 5, 6}};
+        xarray<double> expected_cm = {{3, 1, 2}, {6, 4, 5}};
+        ASSERT_EQ(expected_cm, xt::roll(cm, {1}, {1}));
+
+        // Test 14: Column-major with multi-axis roll
+        xarray<double> expected_cm2 = {{6, 4, 5}, {3, 1, 2}};
+        ASSERT_EQ(expected_cm2, xt::roll(cm, {1, 1}, {0, 1}));
+    }
+
     TEST(xmanipulation, repeat_all_elements_of_axis_0_of_int_array_2_times)
     {
         xarray<int> array = {