Changes

Author: cetagostini

pytensor/link/mlx/dispatch/blockwise.py

@@ -0,0 +1,16 @@
+import mlx.core as mx
+
+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

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

@@ -1,8 +1,8 @@
+import mlx.core as mx
+
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.signal.conv import Conv1d
 
-import mlx.core as mx
-
 
 @mlx_funcify.register(Conv1d)
 def mlx_funcify_Conv1d(op, node, **kwargs):