Toma tu tomate William

Author: cetagostini

pytensor/link/mlx/dispatch/__init__.py

@@ -7,4 +7,4 @@
 import pytensor.link.mlx.dispatch.shape
 import pytensor.link.mlx.dispatch.subtensor
 import pytensor.link.mlx.dispatch.core
-# isort: on
+# 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, Switch
+from pytensor.scalar.basic import AND, OR, Add, Mul, ScalarMaximum, ScalarMinimum
 from pytensor.tensor.elemwise import CAReduce, DimShuffle
 from pytensor.tensor.special import Softmax, SoftmaxGrad
 

pytensor/link/mlx/dispatch/math.py

@@ -1,11 +1,27 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch import mlx_funcify
-from pytensor.scalar.basic import Add, Cos, Exp, Log, Mul, Sin, Sub
+from pytensor.scalar.basic import (
+    EQ,
+    GE,
+    GT,
+    LE,
+    LT,
+    NEQ,
+    Add,
+    Cos,
+    Exp,
+    Log,
+    Mul,
+    Pow,
+    Sin,
+    Sub,
+    Switch,
+    TrueDiv,
+)
+from pytensor.scalar.math import Sigmoid
 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)
@@ -19,6 +35,7 @@ def dot(x, y):
 @mlx_funcify.register(Elemwise)
 def mlx_funcify_Elemwise(op, **kwargs):
     if isinstance(op.scalar_op, Add):
+
         def add(*args):
             result = args[0]
             for arg in args[1:]:
@@ -33,6 +50,7 @@ def sub(x, y):
 
         return sub
     elif isinstance(op.scalar_op, Mul):
+
         def mul(*args):
             result = args[0]
             for arg in args[1:]:
@@ -65,39 +83,64 @@ def 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
+    elif isinstance(op.scalar_op, Switch):
+
+        def switch(cond, x, y):
+            return mx.where(cond, x, y)
+
+        return switch
+    elif isinstance(op.scalar_op, Pow):
+
+        def pow(x, y):
+            return mx.power(x, y)
+
+        return pow
+    elif isinstance(op.scalar_op, TrueDiv):
+
+        def true_div(x, y):
+            return mx.divide(x, y)
+
+        return true_div
     else:
         raise NotImplementedError(f"MLX does not support {op.scalar_op}")

pytensor/link/mlx/dispatch/shape.py

@@ -1,5 +1,6 @@
-from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.tensor.shape import SpecifyShape
+
 
 @mlx_funcify.register(SpecifyShape)
 def mlx_funcify_SpecifyShape(op, node, **kwargs):
@@ -12,4 +13,4 @@ def specifyshape(x, *shape):
                 raise ValueError(f"Invalid shape: Expected {shape} but got {x.shape}")
         return x
 
-    return specifyshape
+    return specifyshape