A lot of new code

Author: cetagostini

pytensor/link/mlx/dispatch/__init__.py

@@ -9,4 +9,5 @@
 import pytensor.link.mlx.dispatch.core
 import pytensor.link.mlx.dispatch.signal
 import pytensor.link.mlx.dispatch.signal.conv
+import pytensor.link.mlx.dispatch.blockwise
 # isort: on

pytensor/link/mlx/dispatch/blockwise.py

@@ -1,16 +1,66 @@
 import mlx.core as mx
 
+from pytensor.graph import FunctionGraph
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.blockwise import Blockwise
 
+
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
-    core_f = mlx_funcify(op.core_op)
-    batched_f = core_f
-    for _ in range(op.batch_ndim(node)):
-        batched_f = mx.vmap(batched_f)
-        
-    def wrapped_blockwise_f(*inputs):
-        return batched_f(*inputs)
-
-    return wrapped_blockwise_f
+    # Create a function graph for the core operation
+    core_node = op._create_dummy_core_node(node.inputs)
+    core_fgraph = FunctionGraph(inputs=core_node.inputs, outputs=core_node.outputs)
+
+    # Convert the core function graph to an MLX function
+    tuple_core_fn = mlx_funcify(core_fgraph, **kwargs)
+
+    # If there's only one output, unwrap it from the tuple
+    if len(node.outputs) == 1:
+
+        def core_fn(*inputs):
+            return tuple_core_fn(*inputs)[0]
+    else:
+        core_fn = tuple_core_fn
+
+    # Apply vmap for each batch dimension
+    batch_ndims = op.batch_ndim(node)
+    vmap_fn = core_fn
+    for _ in range(batch_ndims):
+        vmap_fn = mx.vmap(vmap_fn)
+
+    def blockwise_fn(*inputs):
+        # Check for runtime broadcasting compatibility
+        op._check_runtime_broadcast(node, inputs)
+
+        # Handle broadcasting for batched dimensions
+        if batch_ndims > 0:
+            # Get batch shapes for broadcasting
+            batch_shapes = [inp.shape[:batch_ndims] for inp in inputs]
+
+            # Calculate the broadcasted batch shape
+            from functools import reduce
+
+            def broadcast_shapes(shape1, shape2):
+                return tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2, strict=True))
+
+            if batch_shapes:
+                broadcasted_shape = reduce(broadcast_shapes, batch_shapes)
+
+                # Broadcast inputs to the common batch shape
+                broadcasted_inputs = []
+                for inp in inputs:
+                    if inp.shape[:batch_ndims] != broadcasted_shape:
+                        # Create the full target shape
+                        target_shape = broadcasted_shape + inp.shape[batch_ndims:]
+                        # Broadcast the input
+                        broadcasted_inputs.append(mx.broadcast_to(inp, target_shape))
+                    else:
+                        broadcasted_inputs.append(inp)
+
+                # Apply the vectorized function to the broadcasted inputs
+                return vmap_fn(*broadcasted_inputs)
+
+        # No broadcasting needed
+        return vmap_fn(*inputs)
+
+    return blockwise_fn

pytensor/link/mlx/dispatch/elemwise.py

@@ -1,6 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.scalar import Softplus
 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
@@ -59,9 +60,8 @@ def min(x):
             return mx.min(x, axis=op.axis)
 
         return min
-
     else:
-        raise NotImplementedError(f"MLX does not support {op.scalar_op}")
+        raise NotImplementedError(f"MLX does not support Elemwise {op.scalar_op}")
 
 
 @mlx_funcify.register(Softmax)
@@ -83,3 +83,23 @@ def softmax_grad(dy, sm):
         return dy_times_sm - mx.sum(dy_times_sm, axis=axis, keepdims=True) * sm
 
     return softmax_grad
+
+
+@mlx_funcify.register(Softplus)
+def mlx_funcify_Softplus(op, **kwargs):
+    def softplus(x):
+        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

pytensor/link/mlx/dispatch/math.py

@@ -8,21 +8,29 @@
     LE,
     LT,
     NEQ,
+    Abs,
     Add,
     Cos,
     Exp,
     Log,
     Mul,
     Pow,
     Sin,
+    Sqr,
+    Sqrt,
     Sub,
     Switch,
     TrueDiv,
+    Neg,
+    AND,
+    OR,
+    ScalarMaximum,
+    ScalarMinimum,
 )
 from pytensor.scalar.math import Sigmoid
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.math import Dot
-
+from pytensor.scalar import Softplus
 
 @mlx_funcify.register(Dot)
 def mlx_funcify_Dot(op, **kwargs):
@@ -142,5 +150,70 @@ def true_div(x, y):
             return mx.divide(x, y)
 
         return true_div
+    elif isinstance(op.scalar_op, Sqr):
+
+        def sqr(x):
+            return mx.square(x)
+
+        return sqr
+    elif isinstance(op.scalar_op, Sqrt):
+
+        def sqrt(x):
+            return mx.sqrt(x)
+
+        return sqrt
+    elif isinstance(op.scalar_op, Abs):
+
+        def abs(x):
+            return mx.abs(x)
+
+        return abs
+    elif isinstance(op.scalar_op, Softplus):
+        def softplus(x):
+            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
+    elif isinstance(op.scalar_op, Neg):
+
+        def neg(x):
+            return mx.negative(x)
+
+        return neg
+    elif isinstance(op.scalar_op, AND):
+
+        def all(x):
+            return mx.all(x, axis=op.axis)
+
+        return all
+    elif isinstance(op.scalar_op, OR):
+
+        def any(x):
+            return mx.any(x, axis=op.axis)
+
+        return any
+    elif isinstance(op.scalar_op, ScalarMaximum):
+
+        def max(x):
+            return mx.max(x, axis=op.axis)
+
+        return max
+    elif isinstance(op.scalar_op, ScalarMinimum):
+
+        def min(x):
+            return mx.min(x, axis=op.axis)
+
+        return min
     else:
         raise NotImplementedError(f"MLX does not support {op.scalar_op}")

pytensor/link/mlx/dispatch/signal/conv.py

@@ -5,7 +5,7 @@
 
 
 @mlx_funcify.register(Conv1d)
-def mlx_funcify_Conv1d(op, node, **kwargs):
+def mlx_funcify_Conv1d(op, node=None, **kwargs):
     mode = op.mode
 
     def conv1d(data, kernel):