Repository navigation
fix(tensor): return fewer chunks instead of panicking in chunk - #6263
Merged
Merged
Conversation
`chunk` built `chunks - 1` chunks of `ceil(size / chunks)` and gave the last one `size % chunk_size` elements. When fewer chunks are enough to cover the dimension, that last `narrow` is empty (6 elements into 4 chunks) or runs past the end (7 into 5), and both panic. Use `ceil(size / chunks)` as the chunk size and let `split` produce the chunks, so the last chunk holds the rest and fewer chunks are returned, as the documentation already says.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #6263 +/- ##
=======================================
Coverage 66.63% 66.63%
=======================================
Files 875 875
Lines 129889 129903 +14
=======================================
+ Hits 86547 86561 +14
Misses 43342 43342 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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.
Checklist
cargo run-checkshas been executed.Related Issues/PRs
None, found while testing edge cases of
chunk.Changes
Tensor::chunkpanics innarrowwhen fewer chunks than requested already cover the dimension:arange(0..6).chunk(4, 0)(the last chunk would be empty) andarange(0..7).chunk(5, 0)(it would run past the end). It builtchunks - 1chunks ofceil(size / chunks)and gave the last onesize % chunk_sizeelements.The docs already say
chunkmay return fewer chunks, and PyTorch gives[2, 2, 2]and[2, 2, 2, 1]here.chunknow usesceil(size / chunks)as the chunk size and delegates tosplit. Inputs that worked before give the same result.Testing
test_chunk_fewer_chunks_without_remainder(6 into 4) andtest_chunk_fewer_chunks_with_remainder(7 into 5): both panic on main, pass with the fix.cargo test --no-default-features --features flex,std --test tensor(Flex) passes.cargo fmt --checkand clippy on burn-tensor are clean. I didn't run the fullcargo run-checksor the GPU backends.AI Usage
Claude Code (an AI agent) found this while probing tensor edge cases, wrote the fix, tests, changelog entry and this description, and opened this PR, mostly autonomously.