Skip to content

fix: keep single-element strided slice region round-tripping - #4446

Open
axiom-of-choice wants to merge 1 commit into
ml-explore:mainfrom
axiom-of-choice:fix/slice-singleton-strided-region
Open

axiom-of-choice wants to merge 1 commit into
ml-explore:mainfrom
axiom-of-choice:fix/slice-singleton-strided-region

Conversation

@axiom-of-choice

Copy link
Copy Markdown
Contributor
  • I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: AI was used to help write the fix and the regression tests; verified locally on Apple Silicon (rebuild + pytest + exhaustive sweep). PR description written by the author.

Closes #4444.

A strided slice that selects exactly one element from a span of two or more stored a (start, end, strides) triple that no longer round-trips, so Slice::vjp and Slice::vmap re-derived a different region: the cotangent was broadcast over the whole half-open span, vmap returned every element in the span, and negative strides zeroed the gradient. normalize_slice collapsed the stride to 1 for the singleton axis but left stop alone, and took stop by value so it could not correct it.

Fix: take stop by reference and, in the singleton branch, narrow it so the stored triple still selects the same single element: positive stride stop = start + 1; negative stride stop = start - 1; start == 0 with negative stride (which cannot be expressed as a single-element negative-stride span) falls back to the equivalent unit-stride span. The forward kernel is unchanged: it reads the computed out_shape, so this only affects the stored triple that vjp/vmap re-derive from. The simplification still fires, so no performance change. No backend kernel touched.

Tests (both fail on main, pass here):

  • test_slice_grads_single_element: positive stride, longer span, negative stride, and start == 0 negative fallback.
    Each checks the cotangent lands at the position the forward pass read it from
  • test_vmap_strided_slice_single_element: vmap of a single-element strided slice returns the batched version of the un-batched slice, positive and negative stride

Also exhaustively verified over lengths 2..8 strides 1..4 starts, both axes: 140/140 gradients match numpy, forward values unchanged on every combination (CPU backend)

pre-commit run clean. Full python suite: 854 passed, 1 pre-existing failure (test_fft_grads, torch/dlpack Unsupported device_type, also fails on main without this change)

@zcbenz zcbenz added low priority await verification This pull request is non-trivial and requires a human expert to verify its correctness. labels Sep 2, 2026
@axiom-of-choice
axiom-of-choice marked this pull request as ready for review September 11, 2026 11:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness. low priority

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Silently wrong grad and vmap for a strided slice that selects exactly one element

2 participants