dot product of vectors in multidimentional matrices (python, numpy)

Question

I have two matrices A, B, NxKxD dimensions and I want get matrix C, NxKxDxD dimensions, where C[n, k] = A[n, k] x B[n, k].T (here "x" means product of matrices of dimensions Dx1 and 1xD, so the result must be DxD dimensional), so now my code looking like this (here A = B = X):

def square(X):
    out = np.zeros((N, K, D, D))
    for n in range(N):
        for k in range(K):
            out[n, k] = np.dot(X[n, k, :, np.newaxis], X[n, k, np.newaxis, :])
    return out

It may be slow for big N and K because of python's for cycle. Is there some way to make this multiplication in one numpy function?

Divakar Divakar · Accepted Answer · 2016-12-10T11:47:14

It seems you are not using np.dot for sum-reduction, but just for expansion that results in broadcasting. So, you can simply extend the array to have one more dimension with the use of np.newaxis/None and let the implicit broadcasting help out.

Thus, an implementation would be -

X[...,None]*X[...,None,:]

More info on broadcasting specifically how to add new axes could be found in this other post.

dot product of vectors in multidimentional matrices (python, numpy)

1 Answers