Q:

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)

A:

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

