CODEWITHSUNDEEP
pythonnumpy
betelguese05
betelguese05

Reputation: 23

Vectorizing row-wise numpy operation

I'm writing a neural net from scratch and need to implement the following operation: For each row of matrix dY, take the outer product of the same row of another matrix S (same shape as Y) with itself, and multiply that row of dY by that matrix outer(S[i,:], S[i,:]). Also multiply dY * S element-wise and add that to it.

The code below does this, but it's not vectorized. Can you help me speed this up?

out = dY.copy()
for i in range(dY.shape[0]):
    out[i, :] = dY[i, :] * S[i, :] - dY[i, :].dot(np.outer(S[i, :], S[i, :]))

Update: The following takes an (n,m) matrix S and returns a matrix of shape (n,m,m) where for each row, we take the outer product with itself.

np.einsum("ab,ad->abd", S, S)

Update 2: Finally solved it using two applications of np.einsum.

S_outer = np.einsum("ab,ad->abd", S, S)
return dY * S - np.einsum("ab,abc->ac", dY, S_outer)

Upvotes: 1

Views: 106

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114230

Let's break this into parts. You have

out[i, :] = dY[i, :] * S[i, :] - dY[i, :].dot(np.outer(S[i, :], S[i, :]))

Let's write it as

p = dY[i, :] * S[i, :]
si = S[i, :]
q = dY[i, :].dot(si[:, None] * si)
out[i, :] = p - q

Clearly, p can be factored out of the loop entirely ad dY * S. You can compute q by getting a stack of outer products shaped (S.shape[0], S.shape[1], S.shape[1]) and applying @, a.k.a np.matmul. The key is that matmul broadcasts all but the last two dimensions together, while dot effectively takes the outer product. That allows you to specify that dY is a stack of vectors rather than a matrix simply by introducing a new dimension:

out = dY * S - dY[:, None, :] @ (S[:, :, None] * S[:, None, :])

Upvotes: 0

betelguese05
betelguese05

Reputation: 23

Posting the solution I found as an answer as well.

You can do it with two calls to np.einsum(). 

S_outer = np.einsum("ab,ad->abd", S, S)
return dY * S - np.einsum("ab,abc->ac", dY, S_outer)

Upvotes: 1

Related Questions