Skip to content

Scan primitive lacks output_shapes for shapeless compile #4460

Closed

Description

@justinchuby

Problem

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:

[Primitive::output_shapes] CumSum cannot infer output shapes.

The same primitive backs CumProd, CumMin, CumMax, and CumLogAddExp, so they appear to have the same limitation.

Minimal reproduction

from functools import partial
import mlx.core as mx

@partial(mx.compile, shapeless=True)
def cumulative_sum(x):
    return mx.cumsum(x, axis=0)

y = cumulative_sum(mx.arange(4, dtype=mx.float32))
mx.eval(y)

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

  1. added a commit that references this issue on Sep 4, 2026
  2. added a commit that references this issue on Sep 5, 2026
    bb79f84
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions