ml-explore / ml-explore/mlx

Scan primitive lacks output_shapes for shapeless compile

Open
#4,460 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement low priority
Dominant language
C++
Stars
28.5k
Forks
2.3k
Avg merge
3d 8h
Merged PRs (30d)
62

Description

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the Scan primitive's output_shapes entry point and trace the shapeless shape-inference path used by mx.compile. Verify the existing mx.cumsum reproduction, then add coverage for at least two input shapes and confirm cumulative operations compile and preserve the input shape.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.