Skip to content

Fix crash folding an exclusive cumsum - #2760

Merged
TobyRoseman merged 1 commit into
apple:mainfrom
winklemad:fix-cumsum-exclusive-fold
Jul 20, 2026
Merged

Fix crash folding an exclusive cumsum#2760
TobyRoseman merged 1 commit into
apple:mainfrom
winklemad:fix-cumsum-exclusive-fold

Conversation

@winklemad

Copy link
Copy Markdown
Contributor

Fixes #2759.

Problem

Constant-folding a cumsum with exclusive=True crashes at graph-construction time:

from coremltools.converters.mil.mil import Builder as mb
import numpy as np

@mb.program(input_specs=[])
def prog():
    return mb.cumsum(x=np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), axis=0, exclusive=True)
# TypeError: Cannot construct a dtype from an array

cumsum.value_inference did:

data = np.concatenate((np.zeros(zero_shape, data)), axis=axis)

which (1) passes the data ndarray as np.zeros' dtype argument → TypeError, and (2) wraps a single array in np.concatenate instead of a (zeros, data) sequence. Beyond the crash, the exclusive prefix sum was never actually computed — the inclusive data was 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.dtype and shift the inclusive cumsum one step along axis (prepend a zero, drop the last element) to produce the exclusive prefix sum out[0] = 0, out[i] = sum(x[:i]). reverse is unaffected — it still flips around the exclusive step.

Tests

Added test_builder_eval_exclusive, covering the folded exclusive cumsum across every axis and reverse value against the numpy ground truth. It raises TypeError before this change and passes after; the existing TestCumSum cases are unchanged. Verified the value-inference path locally (the @ssa_fn eval tests run without the CoreML runtime).

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.
@TobyRoseman

Copy link
Copy Markdown
Collaborator

@TobyRoseman
TobyRoseman merged commit 02c736d into apple:main Jul 20, 2026
@TobyRoseman

Copy link
Copy Markdown
Collaborator

Thanks @winklemad

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.

cumsum with exclusive=True crashes when constant-folded (TypeError: Cannot construct a dtype from an array)

2 participants