SUUUUUUUUU!!!!!! LIFE IS GOING WELL. MLX FOR MEDIA MIX MODELS BAY

A shout out for the fathers of the day!

Co-Authored-By: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>
Co-Authored-By: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Author: 3 people

pytensor/link/mlx/dispatch/blockwise.py

@@ -105,20 +105,24 @@ def batched_conv1d(
 
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, **kwargs):
+    # 1) If it's a Conv1d Blockwise, use the custom implementation
     if isinstance(op.core_op, Conv1d):
         return blockwise_conv1d(op, node, **kwargs)
-    
-    core_f = mlx_funcify(op.core_op, node)
 
-    def blockwise_f(*inputs):
-        return blockwise_f(*inputs)
+    # 2) Otherwise, get the core python function for this Blockwise
     core_node = op._create_dummy_core_node(node.inputs)
-    
     core_f = mlx_funcify(op.core_op, core_node)
-    blockwise_f = core_f
-    for i in range(op.batch_ndim(node)):
-        blockwise_f = mx.vmap(blockwise_f)
 
+    # 3) Determine how many inputs correspond to batch dimensions
+    n_batch = op.batch_ndim(node)
+
+    # 4) Build in_axes: map only the first n_batch args, keep the rest static
+    in_axes = tuple(0 if i < n_batch else None for i in range(len(node.inputs)))
+
+    # 5) Vectorize (vmap) with in_axes
+    blockwise_f = mx.vmap(core_f, in_axes=in_axes)
+
+    # 6) Return the mapped function
     def blockwise_fun(*inputs):
         return blockwise_f(*inputs)
 

pytensor/link/mlx/dispatch/elemwise.py

@@ -43,12 +43,12 @@ def careduce(x):
             axis = list(range(x.ndim))
 
         if acc_dtype is None:
-            acc_dtype = x.dtype.type
+            acc_dtype = x.dtype
 
         if op_nfunc_spec:
             mlx_op = getattr(mx, op_nfunc_spec[0])
             return mlx_op(x, axis=axis)
-            return mlx_op(x, axis=axis).astype(acc_dtype)
+            # return mlx_op(x, axis=axis).astype(acc_dtype)
 
         # The PyTensor `Op` didn't tell us which NumPy equivalent to use (or
         # there isn't one), so we use this fallback approach