Guys, I'm getting sad. We need help yisus!!!!!

Author: cetagostini

pytensor/link/mlx/dispatch/blockwise.py

@@ -4,26 +4,36 @@
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.signal.conv import Conv1d
 
-def blockwise_conv1d(op, node):
+import numpy as np
+
+def blockwise_conv1d(op, node, **kwargs):
     if op.core_op.mode != "valid":
         raise NotImplementedError("Only 'valid' mode is supported for conv1d")
-    batches_ndim = op.batch_ndim(node)
-    if batches_ndim != 1:
-        raise NotImplementedError("Only 1D batches are supported for conv1d")
+    # batches_ndim = op.batch_ndim(node)
+    # if batches_ndim != 1:
+    #     raise NotImplementedError("Only 1D batches are supported for conv1d")
 
-    _, kernel = node.inputs
-    if not all(kernel.type.broadcastable[:batches_ndim]):
-        raise NotImplementedError("Only 1D batches are supported for conv1d")
+    # _, kernel = node.inputs
+    # if not all(kernel.type.broadcastable[:batches_ndim]):
+    #     raise NotImplementedError("Only 1D batches are supported for conv1d")
 
     def inner_f(x, kernel):
-        x_reshaped = x.reshape(-1, x.shape[-1]).T # shape equals to (N, B) -> N Time as batches all together
-        b = x_reshaped.shape[1] #
-        kernel_squeeze = kernel.reshape(-1)
-        f = kernel_squeeze.shape[0] # Number of filters
-        kernel_reshaped = mx.broadcast_to(a=kernel_squeeze[None, :, None], shape=(b, f, b))
+        *bx, t = x.shape
+        *bk, h = kernel.shape
+
+        b = np.broadcast_shapes(bx, bk)
+
+        x = x.reshape(b + (t,))
+        kernel = kernel.reshape(b + (h,))
+
+        x_reshaped = x.reshape(-1, t).T # shape equals to (N, B) -> N Time as batches all together
+        kernel_squeeze = kernel.reshape(-1, h)
+        b_prod = kernel_squeeze.shape[0]
+
+        kernel_reshaped = mx.broadcast_to(a=kernel_squeeze[None, :, None], shape=(b_prod, h, b_prod))
         conv_result = mx.conv1d(x_reshaped[None, :, :], kernel_reshaped, stride=1, padding=0, dilation=1)
         _, conv_shape, _ = conv_result.shape
-        return mx.moveaxis(a=conv_result, source=-1, destination=0).reshape(x.shape[:-1] + (conv_shape,))
+        return mx.moveaxis(a=conv_result, source=-1, destination=0).reshape(b + (conv_shape,))
     return inner_f
 
 @mlx_funcify.register(Blockwise)