Skip to content

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

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

fix: keep single-element strided slice region round-tripping#4446
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
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