feat: add multi-axis support for xt::roll (numpy.roll parity) #2877
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Checklist
Description
This PR extends the
xt::rollfunction to support multiple axes simultaneously, completing the NumPy-compatible API as originally requested.Related: #1766, #1823
The original implementation in #1823 only supported single-axis roll operations. According to NumPy's
rollsignature, bothshiftandaxisparameters can be tuples of ints:This PR adds the missing multi-axis support:
Commits
1.
feat: add multi-axis support for xt::roll with optimized pointer arithmeticdetail::roll_multi()using recursive pointer arithmetic algorithm extended from the existing single-axis implementationxt::roll()overloads: one for containers (std::vector,std::array) and one for C-style arrays2.
docs: add multi-axis roll examples to documentationdocs/source/numpy.rstwith NumPy-to-xtensor comparison for multi-axis rolldocs/source/quickref/manipulation.rstwith usage examples3.
bench: add benchmark for multi-axis rollbenchmark/benchmark_roll.cppwith comprehensive benchmarks comparing sequential single-axis rolls vs. the new multi-axis implementationBuild & Test
📊 Benchmark Results (click to expand)
Test Environment: Intel Core i7-12700H, 16GB RAM, Ubuntu 24.04.3 LTS
Performance Summary
The multi-axis implementation achieves significant speedups over sequential single-axis roll calls:
The performance gain increases with the number of axes being rolled, as the multi-axis version avoids creating intermediate temporary arrays and performs the roll operation in a single pass through the data.