Block Matmul with different column size


I have many matrix w1, w2, w3…wn with shape (k*n1, k*n2, k*n3…k*nn) and x1, x2, x3…xn with shape (n1*m, n2*m, n3*m…nn*m), I want to get w1@x1, w2@x2, w3@x3 repectively.

The resulting matrix is multiple k*m matrices and can be concatenated into a large matrix with shape ((k*n)*m).

I didn’t find any efficient solution to vectorize this, simply multiply the matrix (w1;w2; …;wn) and (x1; x2; …xn)^T will yield a matrix with shape (k*m).

Edit: I’m dealing with graph neural networks, where there are a different number of nodes in each graph thus the n1, n2, n3 may be different.