Parallelize a matmul-like matrix computation in Numpy

Question

Input X in shape (n,n,m,m),

Output Y in shape (n,n), where Y[i,j]=∑_{k=1}^{n}{||X[i,j]-X[i,k]*X[k,j]||}, with * denoting point-wise multiplication.

The silly for loop version is like:

X = np.random.randint(1,10,size=(5,5,3,3))
n, _, m, _ = X.shape
Y = np.zeros((n, n))
for i in range(n):
    for j in range(n):
        cnt = 0.0
        X_ij = X[i, j] # in shape m x m
        for k in range(n):
            X_ikj = X[i, k] * X[k, j] # point-wise, in shape m x m
            cnt += np.sum(np.abs(X_ij - X_ikj))
        Y[i, j] = cnt

However I'd like to use a numpy parallel matrix computation. Exactly Y[i,j]=∑_{k=1}^{n}{||X[i,j]-X[i,k]*X[k,j]||} has a similar form with matmul. So in my view there are basically two points:

• How to matmul only along the first two dimensions, while keeping point-wise multiplication for the last two dimensions?
• matmul already summarize the n-dim vector {X[i,k]*X[k,j]}_{k in [1,n]}. However there is a function applied to each X[i,k]*X[k,j] before they're summarized.

Any possible idea is appreciated! Thanks.

Answer 1

You can use broadcasting but you need to swap the two axes with transpose:

np.random.seed(1)
X = np.random.randint(1,10,size=(5,5,3,3))

# transpose
# so X_t[j,k] == X[k,j]
X_t = X.transpose(1,0,2,3)

# output
# X_t[None,...]*X[:,None] is X[k,j] * X[i,k]
ret = np.abs(X[:,:,None] - X_t[None,...]*X[:,None]).sum((2,3,4))

# check
(ret==Y).all()
# True

Output (ret)

array([[1108, 1078,  709,  825,  752],
       [1163, 1185,  988, 1034,  910],
       [1043,  973,  828,  926,  706],
       [ 908,  927,  800, 1078,  765],
       [ 990,  905,  662,  864,  865]])