Fix reduction functions to respect the stride of the output #4995
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.
Fixes #4974
Consider reduction ops like
torch.sum,torch.prod, etc, of the formreduce_op(output, input, keepdim)Let
output_sizebe the size of the output withkeepdim=False, andoutput_keepdim_sizebe the size of the output withkeepdim=True.Right now, what happens in a reduce op is the following:
we have an
inputand anoutput.outputis always resized tooutput_keepdim_size. Then, the reduction op is performed with that size, and output is finally either squeezed tooutput_sizeor kept atoutput_keepdim_sizedepending on whatkeepdimis.The problem with what currently happens is that if
keepdim=Falseandoutputinitially has sizeoutput_sizethenoutputwill be resized tooutput_keepdim_size, the reduce op will be performed, and then output will be squeezed tooutput_size. This resize is not a no-op and will affectoutput's contiguity.This PR fixes the issue by always unsqueezing
outputtooutput_keepdim_size. This operation preservesoutput's contiguity.This fixes the following operations:
Test Plan
New unit tests