Vectorizing row-wise numpy operation

Question

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)

Is this homework? (If it is, we'll avoid giving you a straight answer and help you find it on your own) — SIGSTACKFAULT
– SIGSTACKFAULT, Commented Apr 7, 2020 at 15:59
It is a homework, yes. The vectorizing is just for speed though, I've already solved the problem correctly above. — betelguese05
– betelguese05, Commented Apr 7, 2020 at 16:11
Nitpick, but consider out = np.empty_like(dY). You don't need to prefill, much less copy the data. — Mad Physicist
– Mad Physicist, Commented Apr 7, 2020 at 16:17
I've posted a solution using just broadcasting. You should post your solution as an answer to. It legitimately stands on its own. I think you may not need two einsums, ill update of I figure out how. — Mad Physicist
– Mad Physicist, Commented Apr 7, 2020 at 16:54
Don't forget to select an answer by clicking on the check mark under it when you get the chance. We both get some points, and your question will no longer be in the unanswered queue. — Mad Physicist
– Mad Physicist, Commented Apr 7, 2020 at 20:39

betelguese05 · Accepted Answer · 2020-04-07 21:54:45Z

1

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)

answered Apr 7, 2020 at 21:54

betelguese05

234 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Mad Physicist · Accepted Answer · 2020-04-07 22:08:25Z

0

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, :])

edited Apr 7, 2020 at 22:08

answered Apr 7, 2020 at 16:48

Mad Physicist

116k29 gold badges201 silver badges291 bronze badges

3 Comments

betelguese05 Over a year ago

Within q = dY[i, :].dot(s[:, None] * s), you mean si and not s, right?

betelguese05 Over a year ago

Thanks for your answer. Btw, I am also a physicist (though not of the mad variety).

Mad Physicist Over a year ago

@betelgeuse05. Yes. I fixed the typo. Unfortunately I sold out years ago and became an engineer :)

Collectives™ on Stack Overflow

Vectorizing row-wise numpy operation

2 Answers 2

Comments

3 Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

3 Comments

Related