numpy non canonical dot product in higher dimension

Question

I'm trying to vectorize a loop with NumPy but I'm stuck I have a matrix A of shape (NN,NN) I define the A-dot product by

def scalA(u,v):
    return v.T @ A @ u

Then I have two matrices B and C (B has a shape (N,NN) and C has a shape (K,NN) the loop I'm trying to vectorize is

res = np.zeros((N,K))
for n in range(N):
    for k in range(K):
        res[n,k] = scalA(B[n,:], C[k,:])

I found during my research functions like np.tensordot or np.einsum, but I haven't really understood how they work, and (if I have well understood) tensordot will compute the canonical dot product (that would correspond to A = np.eye(NN) in my case).

Thanks !

Answer 1

np.einsum('ni,ji,kj->nk', B,A,C)

I think this works. I wrote it 'by eye' without testing.

Answer 2

Probably you're looking for this:

def scalA(u,v):
    return u @ A @ v.T

If shape of A is (NN,NN), shape of B is (N,NN), and shape of C is (K,NN), the result of scalA(B,C) has shape (N,K)

If shape of A is (NN,NN), shape of B is (NN,), and shape of C is (NN,), the result of scalA(B,C) is just a scalar.

However, if you're expecting B and C to have even higher dimensionality (greater than 2), this may need further tweaking. (I could not tell from your question whether that's the case)