Changes

Author: cetagostini

pytensor/link/mlx/dispatch/__init__.py

@@ -2,4 +2,9 @@
 from pytensor.link.mlx.dispatch.basic import mlx_funcify, mlx_typify
 
 import pytensor.link.mlx.dispatch.math
-# isort: on
+import pytensor.link.mlx.dispatch.basic
+import pytensor.link.mlx.dispatch.elemwise
+import pytensor.link.mlx.dispatch.shape
+import pytensor.link.mlx.dispatch.subtensor
+import pytensor.link.mlx.dispatch.core
+# isort: on

pytensor/link/mlx/dispatch/elemwise.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
-from pytensor.scalar.basic import AND, OR, Add, Mul, ScalarMaximum, ScalarMinimum
+from pytensor.scalar.basic import AND, OR, Add, Mul, ScalarMaximum, ScalarMinimum, Switch
 from pytensor.tensor.elemwise import CAReduce, DimShuffle
 from pytensor.tensor.special import Softmax, SoftmaxGrad
 

pytensor/link/mlx/dispatch/math.py

@@ -4,6 +4,8 @@
 from pytensor.scalar.basic import Add, Cos, Exp, Log, Mul, Sin, Sub
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.math import Dot
+from pytensor.scalar.math import Sigmoid
+from pytensor.scalar.basic import Add, Mul, Sub, Exp, Log, Sin, Cos, LE, LT, GE, GT, EQ, NEQ
 
 
 @mlx_funcify.register(Dot)
@@ -17,9 +19,11 @@ def dot(x, y):
 @mlx_funcify.register(Elemwise)
 def mlx_funcify_Elemwise(op, **kwargs):
     if isinstance(op.scalar_op, Add):
-
-        def add(x, y):
-            return mx.add(x, y)
+        def add(*args):
+            result = args[0]
+            for arg in args[1:]:
+                result = mx.add(result, arg)
+            return result
 
         return add
     elif isinstance(op.scalar_op, Sub):
@@ -29,9 +33,11 @@ def sub(x, y):
 
         return sub
     elif isinstance(op.scalar_op, Mul):
-
-        def mul(x, y):
-            return mx.multiply(x, y)
+        def mul(*args):
+            result = args[0]
+            for arg in args[1:]:
+                result = mx.multiply(result, arg)
+            return result
 
         return mul
     elif isinstance(op.scalar_op, Exp):
@@ -58,5 +64,40 @@ def cos(x):
             return mx.cos(x)
 
         return cos
+    elif isinstance(op.scalar_op, Sigmoid):
+        def sigmoid(x):
+            return mx.sigmoid(x)
+
+        return sigmoid
+    elif isinstance(op.scalar_op, LE):
+        def le(x, y):
+            return mx.less_equal(x, y)
+
+        return le
+    elif isinstance(op.scalar_op, LT):
+        def lt(x, y):
+            return mx.less(x, y)
+
+        return lt
+    elif isinstance(op.scalar_op, GE):
+        def ge(x, y):
+            return mx.greater_equal(x, y)
+
+        return ge
+    elif isinstance(op.scalar_op, GT):
+        def gt(x, y):
+            return mx.greater(x, y)
+
+        return gt
+    elif isinstance(op.scalar_op, EQ):
+        def eq(x, y):
+            return mx.equal(x, y)
+
+        return eq
+    elif isinstance(op.scalar_op, NEQ):
+        def neq(x, y):
+            return mx.not_equal(x, y)
+
+        return neq
     else:
         raise NotImplementedError(f"MLX does not support {op.scalar_op}")

pytensor/link/mlx/dispatch/subtensor.py

@@ -0,0 +1,110 @@
+from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.tensor.subtensor import (
+    AdvancedIncSubtensor,
+    AdvancedIncSubtensor1,
+    AdvancedSubtensor,
+    AdvancedSubtensor1,
+    IncSubtensor,
+    Subtensor,
+    indices_from_subtensor,
+)
+from pytensor.tensor.type_other import MakeSlice
+
+
+BOOLEAN_MASK_ERROR = """MLX does not support resizing arrays with boolean
+masks. In some cases, however, it is possible to re-express your model
+in a form that MLX can compile:
+
+>>> import pytensor.tensor as pt
+>>> x_pt = pt.vector('x')
+>>> y_pt = x_pt[x_pt > 0].sum()
+
+can be re-expressed as:
+
+>>> import pytensor.tensor as pt
+>>> x_pt = pt.vector('x')
+>>> y_pt = pt.where(x_pt > 0, x_pt, 0).sum()
+"""
+
+DYNAMIC_SLICE_LENGTH_ERROR = """MLX does not support slicing arrays with a dynamic
+slice length.
+"""
+
+
+@mlx_funcify.register(Subtensor)
+@mlx_funcify.register(AdvancedSubtensor)
+@mlx_funcify.register(AdvancedSubtensor1)
+def mlx_funcify_Subtensor(op, node, **kwargs):
+    idx_list = getattr(op, "idx_list", None)
+
+    def subtensor(x, *ilists):
+        indices = indices_from_subtensor(ilists, idx_list)
+        if len(indices) == 1:
+            indices = indices[0]
+
+        return x.__getitem__(indices)
+
+    return subtensor
+
+
+@mlx_funcify.register(IncSubtensor)
+@mlx_funcify.register(AdvancedIncSubtensor1)
+def mlx_funcify_IncSubtensor(op, node, **kwargs):
+    idx_list = getattr(op, "idx_list", None)
+
+    if getattr(op, "set_instead_of_inc", False):
+
+        def mlx_fn(x, indices, y):
+            if not op.inplace:
+                x = x.copy()
+            x[indices] = y
+            return x
+
+    else:
+
+        def mlx_fn(x, indices, y):
+            if not op.inplace:
+                x = x.copy()
+            x[indices] += y
+            return x
+
+    def incsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=idx_list):
+        indices = indices_from_subtensor(ilist, idx_list)
+        if len(indices) == 1:
+            indices = indices[0]
+
+        return mlx_fn(x, indices, y)
+
+    return incsubtensor
+
+
+@mlx_funcify.register(AdvancedIncSubtensor)
+def mlx_funcify_AdvancedIncSubtensor(op, node, **kwargs):
+    if getattr(op, "set_instead_of_inc", False):
+
+        def mlx_fn(x, indices, y):
+            if not op.inplace:
+                x = x.copy()
+            x[indices] = y
+            return x
+
+    else:
+
+        def mlx_fn(x, indices, y):
+            if not op.inplace:
+                x = x.copy()
+            x[indices] += y
+            return x
+
+    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn):
+        return mlx_fn(x, ilist, y)
+
+    return advancedincsubtensor
+
+
+@mlx_funcify.register(MakeSlice)
+def mlx_funcify_MakeSlice(op, **kwargs):
+    def makeslice(*x):
+        return slice(*x)
+
+    return makeslice