fix for all

Author: williambdean

pytensor/link/mlx/dispatch/elemwise.py

@@ -39,7 +39,7 @@ def prod(x):
     elif isinstance(op.scalar_op, AND):
 
         def all(x):
-            return mx.all(a=x, axis=op.axis)
+            return x.all(axis=op.axis)
 
         return all
     elif isinstance(op.scalar_op, OR):

pytensor/link/mlx/dispatch/math.py

@@ -17,6 +17,7 @@
     Cos,
     Exp,
     Log,
+    Log1p,
     Mul,
     Neg,
     Pow,
@@ -29,7 +30,6 @@
     Sub,
     Switch,
     TrueDiv,
-    Log1p
 )
 from pytensor.scalar.math import Sigmoid
 from pytensor.tensor.elemwise import Elemwise
@@ -199,7 +199,7 @@ def neg(x):
     elif isinstance(op.scalar_op, AND):
 
         def all(x):
-            return mx.all(a=x, axis=op.axis)
+            return x.all(axis=op.axis)
 
         return all
     elif isinstance(op.scalar_op, OR):

tests/link/mlx/test_elemwise.py

@@ -0,0 +1,12 @@
+import pytensor.tensor as pt
+from tests.link.mlx.test_basic import compare_mlx_and_py, mx
+
+
+def test_all() -> None:
+    x = pt.vector("x")
+
+    out = pt.all(x > 0)
+
+    x_test = mx.array([-1.0, 2.0, 3.0])
+
+    compare_mlx_and_py([x], out, [x_test])