fix: keep single-element strided slice region round-tripping - #4446
Draft
axiom-of-choice wants to merge 1 commit into
Draft
fix: keep single-element strided slice region round-tripping#4446axiom-of-choice wants to merge 1 commit into
axiom-of-choice wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
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, soSlice::vjpandSlice::vmapre-derived a different region: the cotangent was broadcast over the whole half-open span,vmapreturned every element in the span, and negative strides zeroed the gradient.normalize_slicecollapsed the stride to 1 for the singleton axis but leftstopalone, and tookstopby value so it could not correct it.Fix: take
stopby reference and, in the singleton branch, narrow it so the stored triple still selects the same single element: positive stridestop = start + 1; negative stridestop = start - 1;start == 0with 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 computedout_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, andstart == 0negative fallback.Each checks the cotangent lands at the position the forward pass read it from
test_vmap_strided_slice_single_element:vmapof a single-element strided slice returns the batched version of the un-batched slice, positive and negative strideAlso 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 runclean. Full python suite: 854 passed, 1 pre-existing failure (test_fft_grads, torch/dlpackUnsupported device_type, also fails onmainwithout this change)