MLX backend POC

pytensor/link/mlx/dispatch/core.py

@@ -177,7 +177,7 @@ def tensor_from_scalar(x):
 @mlx_funcify.register(ScalarFromTensor)
 def mlx_funcify_ScalarFromTensor(op, **kwargs):
     def scalar_from_tensor(x):
-        return x.reshape(-1)[0]
+        return mx.array(x).reshape(-1)[0]
 
     return scalar_from_tensor
 

pytensor/link/mlx/dispatch/math.py

@@ -18,7 +18,9 @@
     Cast,
     Cos,
     Exp,
+    IntDiv,
     Invert,
+    IsNan,
     Log,
     Log1p,
     Mul,
@@ -34,7 +36,7 @@
     Switch,
     TrueDiv,
 )
-from pytensor.scalar.math import Sigmoid
+from pytensor.scalar.math import Erfc, Erfcx, Sigmoid, Softplus
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.math import Dot
 
@@ -113,6 +115,14 @@ def true_div(x, y):
     return true_div
 
 
+@mlx_funcify_Elemwise_scalar_op.register(IntDiv)
+def _(scalar_op):
+    def int_div(x, y):
+        return mx.floor_divide(x, y)
+
+    return int_div
+
+
 @mlx_funcify_Elemwise_scalar_op.register(Pow)
 def _(scalar_op):
     def pow(x, y):
@@ -309,11 +319,51 @@ def sigmoid(x):
 @mlx_funcify_Elemwise_scalar_op.register(Invert)
 def _(scalar_op):
     def invert(x):
-        return ~x
+        return mx.bitwise_invert(x)
 
     return invert
 
 
+@mlx_funcify_Elemwise_scalar_op.register(IsNan)
+def _(scalar_op):
+    def isnan(x):
+        return mx.isnan(x)
+
+    return isnan
+
+
+@mlx_funcify_Elemwise_scalar_op.register(Erfc)
+def _(scalar_op):
+    def erfc(x):
+        return 1.0 - mx.erf(x)
+
+    return erfc
+
+
+@mlx_funcify_Elemwise_scalar_op.register(Erfcx)
+def _(scalar_op):
+    def erfcx(x):
+        return mx.exp(x * x) * (1.0 - mx.erf(x))
+
+    return erfcx
+
+
+@mlx_funcify_Elemwise_scalar_op.register(Softplus)
+def _(scalar_op):
+    def softplus(x):
+        # Numerically stable implementation of log(1 + exp(x))
+        # Following the same logic as the original PyTensor implementation
+        return mx.where(
+            x < -37.0,
+            mx.exp(x),
+            mx.where(
+                x < 18.0, mx.log1p(mx.exp(x)), mx.where(x < 33.3, x + mx.exp(-x), x)
+            ),
+        )
+
+    return softplus
+
+
 @mlx_funcify.register(Elemwise)
 def mlx_funcify_Elemwise(op, node, **kwargs):
     # Dispatch to the appropriate scalar op handler

pytensor/link/mlx/dispatch/shape.py

@@ -1,5 +1,7 @@
+import mlx.core as mx
+
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
-from pytensor.tensor.shape import Shape, Shape_i, SpecifyShape
+from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 
 
 @mlx_funcify.register(Shape)
@@ -30,3 +32,11 @@ def shape_i(x):
         return x.shape[op.i]
 
     return shape_i
+
+
+@mlx_funcify.register(Reshape)
+def mlx_funcify_Reshape(op, **kwargs):
+    def reshape(x, shp):
+        return mx.reshape(x, shp)
+
+    return reshape

tests/link/mlx/test_basic.py

@@ -1,18 +1,19 @@
 from collections.abc import Callable, Iterable
 from functools import partial
 
+import mlx.core as mx
 import numpy as np
-import pytest
 
+import pytensor
+from pytensor import tensor as pt
 from pytensor.compile.function import function
 from pytensor.compile.mode import MLX, Mode
 from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Variable
 from pytensor.link.mlx import MLXLinker
+from pytensor.link.mlx.dispatch.core import mlx_funcify_ScalarFromTensor
 
 
-mx = pytest.importorskip("mlx.core")
-
 optimizer = RewriteDatabaseQuery(include=["mlx"], exclude=MLX._optimizer.exclude)
 mlx_mode = Mode(linker=MLXLinker(), optimizer=optimizer)
 py_mode = Mode(linker="py", optimizer=None)
@@ -78,3 +79,52 @@ def compare_mlx_and_py(
         assert_fn(mlx_res, py_res)
 
     return pytensor_mlx_fn, mlx_res
+
+
+def test_scalar_from_tensor_with_scalars():
+    """Test ScalarFromTensor works with both MLX arrays and Python/NumPy scalars.
+
+    This addresses the AttributeError that occurred when Python integers were
+    passed to ScalarFromTensor instead of MLX arrays.
+    """
+    scalar_from_tensor_func = mlx_funcify_ScalarFromTensor(None)
+
+    # Test with MLX array
+    mlx_array = mx.array([42])
+    result = scalar_from_tensor_func(mlx_array)
+    assert result == 42
+
+    # Test with Python int (this used to fail)
+    python_int = 42
+    result = scalar_from_tensor_func(python_int)
+    assert result == 42
+
+    # Test with Python float
+    python_float = 3.14
+    result = scalar_from_tensor_func(python_float)
+    assert abs(result - 3.14) < 1e-6
+
+    # Test with NumPy scalar
+    numpy_scalar = np.int32(123)
+    result = scalar_from_tensor_func(numpy_scalar)
+    assert result == 123
+
+    # Test with NumPy float scalar
+    numpy_float = np.float32(2.71)
+    result = scalar_from_tensor_func(numpy_float)
+    assert abs(result - 2.71) < 1e-6
+
+
+def test_scalar_from_tensor_pytensor_integration():
+    """Test ScalarFromTensor in a PyTensor graph context."""
+    # Create a 0-d tensor (scalar tensor)
+    x = pt.as_tensor_variable(42)
+
+    # Apply ScalarFromTensor
+    scalar_result = pt.scalar_from_tensor(x)
+
+    # Create function and test
+    f = pytensor.function([], scalar_result, mode="MLX")
+    result = f()
+
+    assert result == 42

tests/link/mlx/test_elemwise.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pytest
 
 import pytensor.tensor as pt
@@ -11,3 +12,39 @@ def test_input(op) -> None:
     x_test = mx.array([1.0, 2.0, 3.0])
 
     compare_mlx_and_py([x], out, [x_test])
+
+
+def test_new_elemwise_operations() -> None:
+    """Test new elemwise operations (IntDiv, IsNan, Erfc, Erfcx, Softplus) in elemwise context"""
+    x = pt.vector("x")
+    y = pt.vector("y")
+
+    # Test int_div in an elemwise expression
+    out_int_div = pt.int_div(x, y) + 1
+    x_test = mx.array([10.0, 15.0, 20.0])
+    y_test = mx.array([3.0, 4.0, 6.0])
+    compare_mlx_and_py([x, y], out_int_div, [x_test, y_test])
+
+    # Test isnan in an elemwise expression
+    z = pt.vector("z")
+    out_isnan = pt.isnan(z).astype("float32") * 10
+    z_test = mx.array([1.0, np.nan, 3.0])
+    compare_mlx_and_py([z], out_isnan, [z_test])
+
+    # Test erfc in an elemwise expression
+    w = pt.vector("w")
+    out_erfc = pt.erfc(w) * 2.0
+    w_test = mx.array([0.0, 0.5, 1.0])
+    compare_mlx_and_py([w], out_erfc, [w_test])
+
+    # Test erfcx in an elemwise expression
+    v = pt.vector("v")
+    out_erfcx = pt.erfcx(v) + 0.1
+    v_test = mx.array([0.0, 1.0, 2.0])
+    compare_mlx_and_py([v], out_erfcx, [v_test])
+
+    # Test softplus in an elemwise expression
+    u = pt.vector("u")
+    out_softplus = pt.softplus(u) - 0.5
+    u_test = mx.array([0.0, 1.0, -1.0])
+    compare_mlx_and_py([u], out_softplus, [u_test])

tests/link/mlx/test_math.py

@@ -79,6 +79,7 @@ def test_input(op) -> None:
         pytest.param(pt.eq, id="eq"),
         pytest.param(pt.neq, id="neq"),
         pytest.param(pt.true_div, id="true_div"),
+        pytest.param(pt.int_div, id="int_div"),
     ],
 )
 def test_elemwise_two_inputs(op) -> None:
@@ -90,6 +91,119 @@ def test_elemwise_two_inputs(op) -> None:
     compare_mlx_and_py([x, y], out, [x_test, y_test])
 
 
+def test_int_div_specific() -> None:
+    """Test integer division with specific test cases"""
+    x = pt.vector("x")
+    y = pt.vector("y")
+    out = pt.int_div(x, y)
+
+    # Test with integers that demonstrate floor division behavior
+    x_test = mx.array([7.0, 8.0, 9.0, -7.0, -8.0])
+    y_test = mx.array([3.0, 3.0, 3.0, 3.0, 3.0])
+
+    compare_mlx_and_py([x, y], out, [x_test, y_test])
+
+
+def test_isnan() -> None:
+    """Test IsNan operation with various inputs including NaN values"""
+    x = pt.vector("x")
+    out = pt.isnan(x)
+
+    # Test with mix of normal values, NaN, and infinity
+    x_test = mx.array([1.0, np.nan, 3.0, np.inf, -np.nan, 0.0, -np.inf])
+
+    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_isnan_edge_cases() -> None:
+    """Test IsNan with edge cases"""
+    x = pt.scalar("x")
+    out = pt.isnan(x)
+
+    # Test individual cases
+    test_cases = [0.0, np.nan, np.inf, -np.inf, 1e-10, 1e10]
+
+    for test_val in test_cases:
+        x_test = test_val
+        compare_mlx_and_py([x], out, [x_test])
+
+
+def test_erfc() -> None:
+    """Test complementary error function"""
+    x = pt.vector("x")
+    out = pt.erfc(x)
+
+    # Test with various values including negative, positive, and zero
+    x_test = mx.array([0.0, 0.5, 1.0, -0.5, -1.0, 2.0, -2.0, 0.1])
+
+    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_erfc_extreme_values() -> None:
+    """Test erfc with extreme values"""
+    x = pt.vector("x")
+    out = pt.erfc(x)
+
+    # Test with larger values where erfc approaches 0 or 2
+    x_test = mx.array([-3.0, -2.5, 2.5, 3.0])
+
+    # Use relaxed tolerance for extreme values due to numerical precision differences
+    from functools import partial
+
+    relaxed_assert = partial(np.testing.assert_allclose, rtol=1e-3, atol=1e-6)
+
+    compare_mlx_and_py([x], out, [x_test], assert_fn=relaxed_assert)
+
+
+def test_erfcx() -> None:
+    """Test scaled complementary error function"""
+    x = pt.vector("x")
+    out = pt.erfcx(x)
+
+    # Test with positive values where erfcx is most numerically stable
+    x_test = mx.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5])
+
+    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_erfcx_small_values() -> None:
+    """Test erfcx with small values"""
+    x = pt.vector("x")
+    out = pt.erfcx(x)
+
+    # Test with small values
+    x_test = mx.array([0.001, 0.01, 0.1, 0.2])
+
+    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_softplus() -> None:
+    """Test softplus (log(1 + exp(x))) function"""
+    x = pt.vector("x")
+    out = pt.softplus(x)
+
+    # Test with normal range values
+    x_test = mx.array([0.0, 1.0, 2.0, -1.0, -2.0, 10.0])
+
+    compare_mlx_and_py([x], out, [x_test])
+
+
+def test_softplus_extreme_values() -> None:
+    """Test softplus with extreme values to verify numerical stability"""
+    x = pt.vector("x")
+    out = pt.softplus(x)
+
+    # Test with extreme values where different branches of the implementation are used
+    x_test = mx.array([-40.0, -50.0, 20.0, 30.0, 35.0, 50.0])
+
+    # Use relaxed tolerance for extreme values due to numerical precision differences
+    from functools import partial
+
+    relaxed_assert = partial(np.testing.assert_allclose, rtol=1e-4, atol=1e-8)
+
+    compare_mlx_and_py([x], out, [x_test], assert_fn=relaxed_assert)
+
+
 @pytest.mark.xfail(reason="Argmax not implemented yet")
 def test_mlx_max_and_argmax():
     # Test that a single output of a multi-output `Op` can be used as input to