Fix crash folding an exclusive cumsum - #2760
Merged
Merged
Conversation
cumsum.value_inference passed the data ndarray as np.zeros' dtype argument (np.zeros(zero_shape, data)) and wrapped a single array in np.concatenate, so constant-folding an exclusive=True cumsum raised "TypeError: Cannot construct a dtype from an array" at graph-construction time, and the exclusive prefix sum was never actually computed. Build the leading zeros with the data's dtype and shift the inclusive cumsum one step along the axis (prepend a zero, drop the last element) to produce the exclusive prefix sum. Added value-inference coverage for exclusive cumsum across every axis/reverse combination.
Collaborator
TobyRoseman
reviewed
Jul 20, 2026
TobyRoseman
approved these changes
Jul 20, 2026
Collaborator
|
Thanks @winklemad |
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.
Fixes #2759.
Problem
Constant-folding a
cumsumwithexclusive=Truecrashes at graph-construction time:cumsum.value_inferencedid:which (1) passes the
datandarray asnp.zeros'dtypeargument →TypeError, and (2) wraps a single array innp.concatenateinstead of a(zeros, data)sequence. Beyond the crash, the exclusive prefix sum was never actually computed — the inclusivedatawas returned unshifted. This is reachable from real frontends (e.g.tf.cumsum(..., exclusive=True)on a foldable tensor).Fix
Build the leading zeros with
data.dtypeand shift the inclusive cumsum one step alongaxis(prepend a zero, drop the last element) to produce the exclusive prefix sumout[0] = 0,out[i] = sum(x[:i]).reverseis unaffected — it still flips around the exclusive step.Tests
Added
test_builder_eval_exclusive, covering the folded exclusive cumsum across every axis andreversevalue against the numpy ground truth. It raisesTypeErrorbefore this change and passes after; the existingTestCumSumcases are unchanged. Verified the value-inference path locally (the@ssa_fneval tests run without the CoreML runtime).