Requested changes by Ricardo

Author: cetagostini

pytensor/link/mlx/dispatch/basic.py

@@ -18,14 +18,14 @@ def mlx_typify(data, **kwargs):
 
 
 @mlx_typify.register(np.ndarray)
-@mlx_typify.register(mx.array)
 def mlx_typify_tensor(data, dtype=None, **kwargs):
     return mx.array(data, dtype=dtype)
 
 
 @mlx_typify.register(slice)
 @mlx_typify.register(NoneType)
 @mlx_typify.register(np.number)
+@mlx_typify.register(mx.array)
 def mlx_typify_no_conversion_needed(data, **kwargs):
     return data
 

pytensor/link/mlx/dispatch/core.py

@@ -13,10 +13,10 @@
 
 import warnings
 
-import mlx.core as mx  # MLX
+import mlx.core as mx  
 import numpy as np
 
-from pytensor.link.mlx.dispatch.basic import mlx_funcify  # MLX
+from pytensor.link.mlx.dispatch.basic import mlx_funcify  
 from pytensor.tensor import get_vector_length
 from pytensor.tensor.basic import (
     Alloc,
@@ -34,28 +34,22 @@
 from pytensor.tensor.exceptions import NotScalarConstantError
 
 
-# ------------------------------------------------------------------
-# Join
-# ------------------------------------------------------------------
-@mlx_funcify.register(Join)  # MLX
+@mlx_funcify.register(Join)  
 def mlx_funcify_Join(op, **kwargs):
     def join(axis, *tensors):
         view = op.view
         if (view != -1) and all(
-            tensors[i].shape[axis] == 0  # MLX
+            tensors[i].shape[axis] == 0  
             for i in list(range(view)) + list(range(view + 1, len(tensors)))
         ):
             return tensors[view]
 
-        return mx.concatenate(tensors, axis=axis)  # MLX
+        return mx.concatenate(tensors, axis=axis)  
 
     return join
 
 
-# ------------------------------------------------------------------
-# Split
-# ------------------------------------------------------------------
-@mlx_funcify.register(Split)  # MLX
+@mlx_funcify.register(Split)  
 def mlx_funcify_Split(op: Split, node, **kwargs):
     _, axis_sym, splits_sym = node.inputs
 
@@ -90,7 +84,7 @@ def split(x, axis, splits):
             cumsum_splits = np.cumsum(splits[:-1])
         else:
             # dynamic - keep in graph
-            splits_arr = mx.array(splits)  # MLX
+            splits_arr = mx.array(splits)  
             cumsum_splits = mx.cumsum(
                 splits_arr[:-1]
             ).tolist()  # python list for mx.split
@@ -104,33 +98,29 @@ def split(x, axis, splits):
         if np.any(np.asarray(splits) < 0):
             raise ValueError("Split sizes cannot be negative.")
 
-        return mx.split(x, cumsum_splits, axis=axis)  # MLX
+        return mx.split(x, cumsum_splits, axis=axis)  
 
     return split
 
 
-# ------------------------------------------------------------------
-# ExtractDiag
-# ------------------------------------------------------------------
-@mlx_funcify.register(ExtractDiag)  # MLX
+
+@mlx_funcify.register(ExtractDiag)  
 def mlx_funcify_ExtractDiag(op, **kwargs):
     offset, axis1, axis2 = op.offset, op.axis1, op.axis2
 
     def extract_diag(x, offset=offset, axis1=axis1, axis2=axis2):
-        return mx.diagonal(x, offset=offset, axis1=axis1, axis2=axis2)  # MLX
+        return mx.diagonal(x, offset=offset, axis1=axis1, axis2=axis2)  
 
     return extract_diag
 
 
-# ------------------------------------------------------------------
-# Eye
-# ------------------------------------------------------------------
-@mlx_funcify.register(Eye)  # MLX
+
+@mlx_funcify.register(Eye)  
 def mlx_funcify_Eye(op, **kwargs):
     dtype = convert_dtype_to_mlx(op.dtype)
 
     def eye(N, M, k):
-        return mx.eye(int(N), int(M), int(k), dtype=dtype)  # MLX
+        return mx.eye(int(N), int(M), int(k), dtype=dtype)  
 
     return eye
 
@@ -176,45 +166,37 @@ def convert_dtype_to_mlx(dtype_str):
     return dtype_str
 
 
-# ------------------------------------------------------------------
-# MakeVector
-# ------------------------------------------------------------------
-@mlx_funcify.register(MakeVector)  # MLX
+
+@mlx_funcify.register(MakeVector)  
 def mlx_funcify_MakeVector(op, **kwargs):
     dtype = convert_dtype_to_mlx(op.dtype)
 
     def makevector(*x):
-        return mx.array(x, dtype=dtype)  # MLX
+        return mx.array(x, dtype=dtype)  
 
     return makevector
 
 
-# ------------------------------------------------------------------
-# TensorFromScalar  (identity for MLX)
-# ------------------------------------------------------------------
-@mlx_funcify.register(TensorFromScalar)  # MLX
+
+@mlx_funcify.register(TensorFromScalar)  
 def mlx_funcify_TensorFromScalar(op, **kwargs):
     def tensor_from_scalar(x):
         return x  # already an MLX array / scalar
 
     return tensor_from_scalar
 
 
-# ------------------------------------------------------------------
-# ScalarFromTensor
-# ------------------------------------------------------------------
-@mlx_funcify.register(ScalarFromTensor)  # MLX
+
+@mlx_funcify.register(ScalarFromTensor)  
 def mlx_funcify_ScalarFromTensor(op, **kwargs):
     def scalar_from_tensor(x):
-        return mx.array(x).reshape(-1)[0]  # MLX
+        return mx.array(x).reshape(-1)[0]  
 
     return scalar_from_tensor
 
 
-# ------------------------------------------------------------------
-# Tri
-# ------------------------------------------------------------------
-@mlx_funcify.register(Tri)  # MLX
+
+@mlx_funcify.register(Tri)  
 def mlx_funcify_Tri(op, node, **kwargs):
     # node.inputs  ->  N, M, k
     const_args = [getattr(inp, "data", None) for inp in node.inputs]
@@ -226,7 +208,7 @@ def tri(*args):
             arg if const_a is None else const_a
             for arg, const_a in zip(args, const_args, strict=True)
         ]
-        return mx.tri(*args, dtype=dtype)  # MLX
+        return mx.tri(*args, dtype=dtype)  
 
     return tri
 

pytensor/link/mlx/dispatch/elemwise.py

@@ -1,5 +1,6 @@
 import mlx.core as mx
 import numpy as np
+from functools import singledispatch
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
 from pytensor.link.mlx.dispatch.core import convert_dtype_to_mlx
@@ -10,6 +11,8 @@
     Add,
     Cast,
     Mul,
+    ScalarMaximum,
+    ScalarMinimum,
 )
 from pytensor.tensor.elemwise import CAReduce, DimShuffle
 from pytensor.tensor.special import Softmax, SoftmaxGrad
@@ -32,34 +35,64 @@ def dimshuffle(x):
     return dimshuffle
 
 
-@mlx_funcify.register(CAReduce)
-def mlx_funcify_CAReduce(op, **kwargs):
-    if isinstance(op.scalar_op, Add):
+# Second-level dispatch for scalar operations in CAReduce
+@singledispatch
+def mlx_funcify_CAReduce_scalar_op(scalar_op):
+    raise NotImplementedError(f"MLX does not support CAReduce with scalar op {scalar_op}")
+
+
+@mlx_funcify_CAReduce_scalar_op.register(Add)
+def _(scalar_op):
+    def sum_reduce(x, axis):
+        return mx.sum(x, axis=axis)
+    return sum_reduce
+
+
+@mlx_funcify_CAReduce_scalar_op.register(Mul)
+def _(scalar_op):
+    def prod_reduce(x, axis):
+        return mx.prod(x, axis=axis)
+    return prod_reduce
 
-        def sum(x):
-            return mx.sum(x, axis=op.axis)
 
-        return sum
-    elif isinstance(op.scalar_op, Mul):
+@mlx_funcify_CAReduce_scalar_op.register(AND)
+def _(scalar_op):
+    def all_reduce(x, axis):
+        return x.all(axis=axis)
+    return all_reduce
 
-        def prod(x):
-            return mx.prod(x, axis=op.axis)
 
-        return prod
-    elif isinstance(op.scalar_op, AND):
+@mlx_funcify_CAReduce_scalar_op.register(OR)
+def _(scalar_op):
+    def any_reduce(x, axis):
+        return mx.any(x, axis=axis)
+    return any_reduce
 
-        def all(x):
-            return x.all(axis=op.axis)
 
-        return all
-    elif isinstance(op.scalar_op, OR):
+@mlx_funcify_CAReduce_scalar_op.register(ScalarMaximum)
+def _(scalar_op):
+    def max_reduce(x, axis):
+        return mx.max(x, axis=axis)
+    return max_reduce
 
-        def any(x):
-            return mx.any(x, axis=op.axis)
 
-        return any
-    else:
-        raise NotImplementedError(f"MLX does not support Elemwise {op.scalar_op}")
+@mlx_funcify_CAReduce_scalar_op.register(ScalarMinimum)
+def _(scalar_op):
+    def min_reduce(x, axis):
+        return mx.min(x, axis=axis)
+    return min_reduce
+
+
+@mlx_funcify.register(CAReduce)
+def mlx_funcify_CAReduce(op, **kwargs):
+    # Dispatch to the appropriate scalar op handler
+    scalar_reduce_fn = mlx_funcify_CAReduce_scalar_op(op.scalar_op)
+    axis = op.axis
+    
+    def reduce(x):
+        return scalar_reduce_fn(x, axis)
+    
+    return reduce
 
 
 @mlx_funcify.register(Softmax)