From 2dc7af3206ca719de1f1f27d6e0df137fe9d454d Mon Sep 17 00:00:00 2001 From: keeeeenw Date: Mon, 14 Sep 2026 14:14:35 -0700 Subject: [PATCH] Support shapeless compilation of scan operations --- mlx/primitives.h | 1 + python/tests/test_compile.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/mlx/primitives.h b/mlx/primitives.h index 0cfc71bf04..f5ffc7d86d 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -1888,6 +1888,7 @@ class Scan : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS(); + DEFINE_INPUT_OUTPUT_SHAPE() const char* name() const override { switch (reduce_type_) { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 03d0027d7e..ed8a7369b7 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -7,6 +7,7 @@ import threading from functools import partial, wraps from io import StringIO +from itertools import product import mlx.core as mx import mlx_tests @@ -651,6 +652,33 @@ def fun(x): cfun = mx.compile(fun, shapeless=True) self.assertTrue(mx.array_equal(fun(x2), cfun(x2))) + def test_shapeless_compile_scan(self): + ops = (mx.cumsum, mx.cumprod, mx.cummin, mx.cummax, mx.logcumsumexp) + for op, axis, reverse, inclusive in product( + ops, (0, 1, -1), (False, True), (False, True) + ): + with self.subTest( + op=op.__name__, axis=axis, reverse=reverse, inclusive=inclusive + ): + scan = partial(op, axis=axis, reverse=reverse, inclusive=inclusive) + trace_count = 0 + + def fun(x): + nonlocal trace_count + trace_count += 1 + return scan(x) + + cfun = mx.compile(fun, shapeless=True) + for shape in ((3, 4), (5, 7), (1, 2), (2, 1)): + with self.subTest(shape=shape): + x = mx.arange(math.prod(shape), dtype=mx.float32) + x = ((x % 7 - 3) / 4).reshape(shape) + expected = scan(x) + actual = cfun(x) + self.assertEqual(actual.shape, x.shape) + self.assertEqualArray(actual, expected, atol=1e-6, rtol=1e-6) + self.assertEqual(trace_count, 1) + def test_shapeless_compile_unflatten(self): x = mx.zeros((1, 1, 4 * 32))