Skip to content

Fix out-of-bounds accesses on out-of-range axes - #2951

Open
gouarin wants to merge 1 commit into
xtensor-stack:masterfrom
gouarin:fix-axis-validation
Open

gouarin wants to merge 1 commit into
xtensor-stack:masterfrom
gouarin:fix-axis-validation

Conversation

@gouarin

@gouarin gouarin commented Oct 29, 2026

Copy link
Copy Markdown
Member

Checklist

  • The title and commit message(s) are descriptive.
  • Small commits made to fix your PR have been squashed to avoid history pollution.
  • Tests have been added for new features or bug fixes.
  • API of new functions and classes are documented (no new public API; the normalize_axis doc comment is updated).

Description

normalize_axis did not check its input, and several functions taking an axis did not call it. An out-of-range axis was then used to index shape, strides or slice containers:

xt::xtensor<double, 2> a = {{1., 2.}, {3., 0.}};
xt::argmax(a, 2);  // writes past the std::array<std::size_t, 1> result shape

AddressSanitizer reports buffer overflows of this kind in argmax, diff, expand_dims, stack, axis_slice_begin and fft. The container overloads of normalize_axis only had an XTENSOR_ASSERT, which is off in release builds.

normalize_axis now throws std::runtime_error when the axis is outside [-dim, dim), the type reducers, roll, split and squeeze already throw. It is a template on the axis type, so std::size_t(-1) is rejected instead of becoming the last axis. expand_dims and stack check against dim + 1 since they insert an axis. Functions that skipped the check (flip, repeat, squeeze, concatenate, stack, expand_dims, the axis iterators, the blockwise reducers, and the 1-D shortcut of sort/argsort/partition/argpartition) now call it. In argmin/argmax, the result shape is built with a loop, which also removes the GCC 13/15 -Warray-bounds warnings in xsort.hpp.

Behavior changes:

  • an out-of-range axis throws instead of being undefined behavior;
  • squeeze and the blockwise reducers accept negative axes and reject repeated ones;
  • average with 1-D weights requires exactly one axis;
  • the error message of reducers, roll and split changes, not the exception type.

The tests added in test_xutils, test_xsort, test_xmath, test_xmanipulation, test_xbuilder, test_xaxis_iterator, test_xaxis_slice_iterator, test_xfft and test_xblockwise_reducer fail on master and pass with the fix, also under AddressSanitizer and UBSan.

Functions taking an axis used it to index shape, strides or slice
containers without checking it. xt::argmax(a, 2) on a 2-D xtensor wrote
past the std::array holding the result shape (stack-buffer-overflow
under AddressSanitizer). diff, trapz, unwrap, cumsum, average, sort,
partition, median, quantile, fft, expand_dims, flip, repeat, squeeze,
concatenate, stack, the axis iterators and the blockwise reducers read
or wrote out of bounds the same way.

normalize_axis did no validation, and its container overloads only had
an XTENSOR_ASSERT, which is off in release builds. It also took the
axis as std::ptrdiff_t, so an unsigned std::size_t(-1) became the last
axis instead of an error.

normalize_axis now throws std::runtime_error when the axis is not in
[-dim, dim). That is the exception type reducers, roll, split and
squeeze already throw for a bad axis. It is a template over the axis
type, and negative values count from the end only for signed types.
The container overloads and forward_normalize run the same check.
expand_dims and stack insert a new axis, so they validate against
dim + 1.

Functions that indexed with the raw axis now go through normalize_axis:
expand_dims, split, flip, repeat, squeeze, concatenate, stack,
axis_begin, axis_end, axis_slice_begin, axis_slice_end and the
blockwise reducers. sort, partition and argpartition check the axis
before their 1-D shortcut, which used to ignore it. roll and split drop
their own bounds check in favor of the shared one.

argmin and argmax validate the axis inside arg_func_impl, next to the
code that builds the result shape. That code now uses a loop that skips
the axis instead of two std::copy calls, which removes the GCC
-Warray-bounds warnings on alt_shape.

squeeze and the blockwise reducers also reject repeated axes, because
they size the result shape from the number of axes. squeeze accepts
negative axes. average throws when 1-D weights come with anything other
than exactly one axis, instead of reading ax[0].
@codspeed

codspeed Bot commented Oct 29, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will regress 1 benchmark

2 warnings

⚠️ Unable to measure the execution of 1 benchmark

This benchmark measured no execution time, so its result is not comparable and counts as unchanged. This usually happens when the compiler removes the code under test.

Prevent compiler optimizations

See affected benchmarks

⚠️ Base and head ran on different runner types

Some benchmarks with significant performance changes were compared across different runner types, so the changes might not come from the code. Run base and head on the same runner type, or use CodSpeed Macro Runners for more stable results.

Learn more about Macro Runners

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
❌ 1 regressed benchmark
✅ 253 untouched benchmarks

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

Benchmark BASE HEAD Efficiency
❌ xshape_initializer[std\:\:array<std\:\:size_t, 4>] 151.7 ns 180.8 ns -16.13%
⚡ xshape_access[std\:\:array<std\:\:size_t, 4>] 154.4 ns 125.3 ns +23.28%

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.


Comparing gouarin:fix-axis-validation (b9bf855) with master (57f3d07)

Open in CodSpeed

@gouarin gouarin mentioned this pull request Oct 29, 2026
3 of 4 tasks

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant