mps.matmul error on dimensions differ without any matmul in code

Question

I am writing a MPS Graph code to adapt Denoiser to run on iOS device, I had to implement Conv1d, ConvTranspose1d to use the 2d implemented on MPS Graph. I face an error on the run of the graph :

(mpsFileLoc): /Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:39:0: error: 'mps.matmul' op contracting dimensions differ 1024 & 2048 (mpsFileLoc): /Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:39:0: note: see current operation: %212 = "mps.matmul"(%208#0, %211) <{transpose_lhs = false, transpose_rhs = false}> : (tensor<7x1x1024xf32>, tensor<2048x4096xf32>) -> tensor<7x1x4096xf32>

I do not have any matrix multiplication in my code, so I suspect it might be in a convolution or transposed convolution ? Here print of the shapes of the tensors that are part of any convolution, and my code to transform a 1d convolution to a 2d convolution.

1D->2D CONVTRANS input shape [1, 7, 1, 1024]
1D->2D CONVTRANS weight shape [1, 8, 1024, 512]
1D->2D CONVTRANS input shape [1, 32, 1, 512]
1D->2D CONVTRANS weight shape [1, 8, 512, 256]
1D->2D CONVTRANS input shape [1, 132, 1, 256]
1D->2D CONVTRANS weight shape [1, 8, 256, 128]
1D->2D CONVTRANS input shape [1, 532, 1, 128]
1D->2D CONVTRANS weight shape [1, 8, 128, 64]
1D->2D CONVTRANS input shape [1, 2132, 1, 64]
1D->2D CONVTRANS weight shape [1, 8, 64, 1]

    public func convolutionTranspose1D(source: MPSGraphTensor, weights: MPSGraphTensor, biases: MPSGraphTensor, outputShape: [NSNumber], descriptor: MPSGraphConvolution1dOpDescriptor, name: String? = nil) -> MPSGraphTensor {

        guard source.shape?.count == 3 else { fatalError("Source needs to be dimension of 4") }
        guard weights.shape?.count == 3 else { fatalError("Weights need to be dimension of 4") }
        var source = expandDims(source, axes: [2], name: name)
        source = transpose(source, permutation: [0, 3, 2, 1], name: name)
        print("1D->2D CONVTRANS input shape \(source.shape!)")
        var weights = expandDims(weights, axes: [2], name: name)
        weights = transpose(weights, permutation: [2, 3, 1, 0], name: name)
        print("1D->2D CONVTRANS weight shape \(weights.shape!)")
        var outputShape = outputShape
        outputShape.insert(1, at: 2)
        let conv2DDescriptor = descriptor.descriptor2D
        source = convolutionTranspose2D(source, weights: weights, outputShape: outputShape, descriptor: conv2DDescriptor, name: name)

        source = squeeze(source, axes: [2], name: name)

        return addition(source, biases, name: name)
    }

    private func convolutionTranspose1DLoad(source: MPSGraphTensor, chIn: Int, chOut: Int, padding: (Int, Int) = (0, 0), kernelSize: Int, stride: Int, dilation: Int = 1, path: String, name: String? = nil) -> MPSGraphTensor {

        let convWeights = fileLoad("\(path).weight", shape: [chOut as NSNumber, chIn as NSNumber, kernelSize as NSNumber], name: nil)
        let convBiases = fileLoad("\(path).bias", shape: [1, chOut as NSNumber, 1], name: nil)

        let (B, _, L) = source.shape!.dims3D
        let Lout = (L.intValue - 1) * stride - padding.0 - padding.1 + dilation * (kernelSize - 1) + 1
        
        let outputShape = [B, chOut as NSNumber, Lout as NSNumber]
        return convolutionTranspose1D(source: source, weights: convWeights, biases: convBiases, outputShape: outputShape, descriptor:  MPSGraphConvolution1dOpDescriptor(stride: stride, dilationRate: dilation, paddingLeft: padding.0, paddingRight: padding.1, paddingStyle: .explicit, dataLayout: .NCHW, weightsLayout: .HWIO))
    }

I exported my weights with

    private func fileLoad(_ file: String, shape: [NSNumber], name: String? = nil) -> MPSGraphTensor {
        let modelDataURL = Bundle.main.url(forResource: "\(file)", withExtension: "bin")
        let modelData = try! Data(contentsOf: modelDataURL!)
        return constant(modelData, shape: shape, dataType: .float32)
    }

for name, param in model.named_parameters():
    filename = str(name) + '.bin'
    path = os.path.join(basedir, filename)
    print(path)
    torch.save(param, path)