Handle dynamic shapes to AllocEmpty in non-compiled mode

Author: jessegrabowski

pytensor/link/mlx/dispatch/core.py

@@ -247,18 +247,29 @@ def tri(*args):
 
 
 @mlx_funcify.register(AllocEmpty)
-def mlx_funcify_AllocEmpty(op, **kwargs):
+def mlx_funcify_AllocEmpty(op, node, **kwargs):
     dtype = convert_dtype_to_mlx(op.dtype)
+    node_inputs = node.inputs
+    static_dims = (
+        _extract_static_dims(node_inputs)
+        if node_inputs and len(node_inputs) > 1
+        else None
+    )
 
     def allocempty(*shape):
-        return mx.zeros(shape, dtype=dtype)
+        resolved_shape = (
+            _resolve_shape(static_dims, shape)
+            if static_dims is not None
+            else tuple(_coerce_to_int(dim) for dim in shape)
+        )
+        return mx.zeros(resolved_shape, dtype=dtype)
 
     return allocempty
 
 
 @mlx_funcify.register(Alloc)
 def mlx_funcify_Alloc(op, node, **kwargs):
-    node_inputs = getattr(node, "inputs", None)
+    node_inputs = node.inputs
     static_dims = (
         _extract_static_dims(node_inputs[1:])
         if node_inputs and len(node_inputs) > 1

tests/link/mlx/test_core.py

@@ -73,23 +73,18 @@ def test_alloc_compilation_limitation():
     # Test that it works with concrete values (non-compiled context)
     output = f(5.0, 3, 4)
     assert output.shape == (3, 4)
-    assert np.allclose(output, 5.0)
+    np.testing.assert_allclose(output, 5.0)
 
     # Test that compilation fails with helpful error
     compiled_f = pytensor.function([x, s1, s2], result, mode=compile_mode)
 
-    with pytest.raises(ValueError) as exc_info:
+    with pytest.raises(
+        ValueError,
+        match="MLX compilation limitation: Alloc operations with dynamic shapes cannot be "
+        "used inside compiled functions",
+    ):
         compiled_f(5.0, 3, 4)
 
-    error_msg = str(exc_info.value)
-    assert "MLX compilation limitation" in error_msg
-    assert "Alloc operations with dynamic shapes" in error_msg
-    assert "cannot be used inside compiled functions" in error_msg
-    assert "Workarounds:" in error_msg
-    assert "Avoid using Alloc with dynamic shapes in compiled contexts" in error_msg
-    assert "Use static shapes when possible" in error_msg
-    assert "Move Alloc operations outside compiled functions" in error_msg
-
 
 def test_alloc_static_shapes_compilation():
     """Test that Alloc operations with static shapes work fine in compiled contexts."""
@@ -109,6 +104,36 @@ def test_alloc_static_shapes_compilation():
 
     assert output_normal.shape == (3, 4)
     assert output_compiled.shape == (3, 4)
-    assert np.allclose(output_normal, 5.0)
-    assert np.allclose(output_compiled, 5.0)
-    assert np.allclose(output_normal, output_compiled)
+    np.testing.assert_allclose(output_normal, 5.0)
+    np.testing.assert_allclose(output_compiled, 5.0)
+    np.testing.assert_allclose(output_normal, output_compiled)
+
+
+def test_empty_static_shape():
+    result = pt.empty((3, 4), dtype="float32")
+
+    f = pytensor.function([], result, mode="MLX")
+    output = f()
+
+    assert output.shape == (3, 4)
+    np.testing.assert_allclose(output, 0.0)
+
+
+def test_empty_dynamic_shape():
+    s1 = pt.scalar("s1", dtype="int64")
+    s2 = pt.scalar("s2", dtype="int64")
+    result = pt.empty((s1, s2), dtype="float32")
+
+    f = pytensor.function([s1, s2], result, mode=mlx_mode_no_compile)
+    output = f(3, 4)
+
+    assert output.shape == (3, 4)
+    np.testing.assert_allclose(output, 0.0)
+
+    f_compiled = pytensor.function([s1, s2], result, mode=compile_mode)
+    with pytest.raises(
+        ValueError,
+        match="MLX compilation limitation: Alloc operations with dynamic shapes cannot be "
+        "used inside compiled functions",
+    ):
+        f_compiled(3, 4)