Skip to content

⁠fix(core): validate shape element count in ShapeWithOneHole blanket impl (#3534)⁠#3572

Open
oyoyo4556 wants to merge 1 commit into
huggingface:mainfrom
oyoyo4556:oyoyo4556-patch-1
Open

⁠fix(core): validate shape element count in ShapeWithOneHole blanket impl (#3534)⁠#3572
oyoyo4556 wants to merge 1 commit into
huggingface:mainfrom
oyoyo4556:oyoyo4556-patch-1

Conversation

@oyoyo4556

Copy link
Copy Markdown

Closes #3534

This PR adds the missing shape element count validation in the ShapeWithOneHole blanket implementation for S: Into<Shape>.

Previously, the _el_count parameter was ignored, allowing the creation of tensors with invalid shapes, which led to out-of-bounds reads and panics (especially via the ONNX loader as reported in #3534).

Please feel free to modify or add regression tests if needed!

Add element count validation in into_shape method.
@oyoyo4556

Copy link
Copy Markdown
Author

I encountered a silent data corruption in my own code: when reusing a pre-allocated Vec buffer, wrong values were silently accepted without any error, making the bug extremely hard to detect.

@pjdurden

Copy link
Copy Markdown

hey, author of #3589 here (just closed it in favor of this one). your fix in into_shape is the more central spot and it covers the same from_slice / from_vec paths mine did, so this is the one that should land.

one thing worth checking before merge though. reshape (tensor.rs:2538) also goes through into_shape and then does its own elem count check right after at tensor.rs:2540, returning Error::ShapeMismatchBinaryOp { op: "reshape" }. with this change into_shape bails first, so a bad reshape now returns the new bail! string instead of the structured error, and reshapes own check at 2540 to 2547 is basically dead code now. probably want to keep reshapes structured error or just drop the unreachable check.

also theres no test here yet. i already have two ready from #3589, they fail before the fix and pass after, covering the constructor entry points: from_slice / from_vec rejecting a shape whose element count doesnt match the data, and from_raw_buffer (goes through from_slice so same contract), plus the matching shape and ((), 3) cases still passing.

maintainers usually like the coverage landing with the fix rather than as a separate PR later, so happy to just add them onto this one. if you add me as a collaborator on your fork i can push the tests straight to your branch, or i can open a small PR against your branch and you merge it in, whatever is easier for you. nice catch on the root cause btw.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Tensor::from_slice / from_raw_buffer return Ok when the declared shape does not match the storage element count

2 participants