Support dynamic shape-derived bounds in MLX ARange

pytensor/link/mlx/dispatch/tensor_basic.py

@@ -178,32 +178,47 @@ def alloc(x, *shape):
     return alloc
 
 
-ARANGE_CONCRETE_VALUE_ERROR = (
-    "MLX's arange requires all arguments (start, stop, step) to be concrete "
-    "Python int/float values, not symbolic variables. Unlike NumPy and JAX, "
-    "MLX does not accept array inputs for arange at all."
-    "\n\nAn example of a valid graph:"
-    "\n>>> import pytensor.tensor as pt"
-    "\n>>> pt.arange(1, 10, 2)"
+ARANGE_DATA_DEPENDENT_ERROR = (
+    "MLX cannot build arange with a data-dependent length: the bounds depend on "
+    "runtime array values, so the output shape is unknown at compile time. "
+    "Constant and shape-derived bounds (e.g. pt.arange(x.shape[0])) are supported."
 )
 
 
 @mlx_funcify.register(ARange)
 def mlx_funcify_ARange(op, node, **kwargs):
-    # MLX's arange only accepts Python int/float, not arrays,
-    # so all arguments must be known at graph-construction time.
-    try:
-        start, stop, step = [
-            get_scalar_constant_value(arg).item() for arg in node.inputs
+    dtype = convert_dtype_to_mlx(op.dtype)
+    # mx.arange only accepts Python int/float. Bake constant bounds, and resolve
+    # the rest at runtime: shape-derived bounds are concrete under mx.compile even
+    # when the static shape is unknown (mirrors the JAX dispatch).
+    static_args = [_arange_static_bound(arg) for arg in node.inputs]
+
+    def arange(*args):
+        resolved = [
+            static if static is not None else _arange_runtime_bound(runtime)
+            for static, runtime in zip(static_args, args, strict=True)
         ]
+        return mx.arange(*resolved, dtype=dtype)
+
+    return arange
+
+
+def _arange_static_bound(arg):
+    try:
+        return get_scalar_constant_value(arg).item()
     except NotScalarConstantError:
-        raise NotImplementedError(ARANGE_CONCRETE_VALUE_ERROR)
-    dtype = convert_dtype_to_mlx(op.dtype)
+        return None
 
-    def arange(*_args):
-        return mx.arange(start, stop, step, dtype=dtype)
 
-    return arange
+def _arange_runtime_bound(value):
+    try:
+        return value.item() if hasattr(value, "item") else value
+    except (ValueError, TypeError) as exc:
+        if "[eval] Attempting to eval an array during function transformations" in str(
+            exc
+        ):
+            raise NotImplementedError(ARANGE_DATA_DEPENDENT_ERROR) from exc
+        raise
 
 
 def _extract_static_dims(shape_inputs):

tests/link/mlx/test_tensor_basic.py

@@ -164,3 +164,48 @@ def test_arange():
     out = arange(1, 10, 2)
 
     compare_mlx_and_py([], [out], [])
+
+
+def test_arange_dynamic_shape():
+    # Shape-derived bounds are concrete under mx.compile even when the static
+    # shape is unknown, so a genuinely dynamic length must work (regression: this
+    # used to raise NotImplementedError because of an over-aggressive constant
+    # check). Exercises every position (start/stop/step) being shape-derived, an
+    # offset, and an empty result.
+    x = pt.vector("x")
+    y = pt.vector("y")
+    outs = [
+        arange(x.shape[0]),  # dynamic stop
+        arange(x.shape[0] + 2),  # shape-derived expression
+        arange(x.shape[0], y.shape[0]),  # dynamic start and stop
+        arange(0, y.shape[0], x.shape[0]),  # dynamic step
+        arange(y.shape[0], x.shape[0]),  # start > stop -> empty
+    ]
+    compare_mlx_and_py(
+        [x, y],
+        outs,
+        [np.zeros(3, dtype="float32"), np.zeros(7, dtype="float32")],
+    )
+
+
+def test_arange_dynamic_advanced_index():
+    # The motivating case: a vectorized gather lowers to advanced indexing that
+    # internally builds arange(idx.shape[0]) with a runtime-dynamic length.
+    logp = pt.matrix("logp")
+    targets = pt.lvector("targets")
+    out = logp[arange(targets.shape[0]), targets]
+    compare_mlx_and_py(
+        [logp, targets],
+        [out],
+        [np.arange(12, dtype="float32").reshape(3, 4), np.array([0, 2, 3])],
+    )
+
+
+def test_arange_data_dependent_raises():
+    # A genuinely data-dependent length has a runtime-only output shape, which MLX
+    # cannot compile. This must fail loudly rather than silently misbehave.
+    x = pt.vector("x")
+    out = arange(pt.sum(x > 0).astype("int64"))
+    fn = pytensor.function([x], out, mode=compile_mode)
+    with pytest.raises(NotImplementedError, match="data-dependent length"):
+        fn(np.array([1.0, -1.0, 1.0], dtype="float32"))