Fix reshape copy bug #1253
Merged
Fix reshape copy bug #1253
Conversation
This file contains 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.
Fixes ml-explore/mlx-swift#104 and potentially ml-explore/mlx-examples#642 .
The two first commits in the PR have two different fixes. The later one is what I consider a better choice. The bug is because the code in
Reshapeassumed thatcopycopies arrays of same total size but different shapes. That is not true, it is not checked but copy assumesin.shape() == out.shape(). I did not add this as an assertion but we could... let me know.The fix in the previous PR was just allowing to copy from any shape to any shape assuming that the total number of elements is the same.
The current fix just fixes reshape to allocate and copy in place passing the correct output strides. I also updated the cpu copy to use
collapse_contiguous_dimsso we should get signficantly faster many-dimensional copies on the cpu.