vectorize NumPy triple product on 2D array

Question

I am trying to vectorize the following triple product operation on an N x N array called p below:

for j in range(len(p)):
    for k in range(len(p)):
        for l in range(len(p)):
            h[j, k, l] = p[j, k] * p[k, l] * p[l, j] - p[j, l] * p[l, k] * p[k, j]

I thought numpy.einsum should be of use here, despite that I'm not actually summing over the repeated indices, but I haven't been able to pin it down. Thoughts?

Answer 1

Simply porting over those loop iterators as string notations, we would have an einsum based solution like so -

h = np.einsum('jk,kl,lj->jkl',p,p,p) - np.einsum('jl,lk,kj->jkl',p,p,p)

Being basically an expansion related question (as we are not reducing any axis), we can simply use NumPy broadcasting too by introducing new axes with None/np.newaxis at various places for allowing the expansion, like so -

h = p[...,None]*p*p[:,None,:].T - p[:,None,:]*p.T*p.T[...,None]