You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Scan primitive lacks output_shapes for shapeless compile #4460
mx.cumsum is implemented by the Scan primitive, but Scan does not override Primitive::output_shapes. As a result, a cumulative operation cannot participate in mx.compile(..., shapeless=True), even though every Scan variant preserves its input shape.
This surfaced in an ONNX Runtime MLX execution provider decoder graph: compiling a shape-polymorphic subgraph containing CumSum reaches shape inference and fails with:
Observed with MLX 0.32.1. The current main branch also has no output_shapes override on Scan.
Expected behavior
The function compiles in shapeless mode and returns an output with the same shape as x, including when the input extent changes between calls.
Suggested fix
Scan::output_shapes can return the first input's shape, equivalent to DEFINE_INPUT_OUTPUT_SHAPE(). A shapeless compile test covering at least two input shapes would prevent regression.
Activity
Use native MLX CumProd safely (#55)
compile: Support shapeless compile for Scan primitive