Векторизация тройного произведения NumPy на двумерном массиве

Question

Я пытаюсь векторизовать следующую тройную операцию продукта на N x N массив называется p ниже:

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]

я думал numpy.einsum должен быть полезен здесь, несмотря на то, что я на самом деле не суммирую по повторяющимся индексам, но я не смог определить это. Мысли?

Answer 1

Если просто перенести эти итераторы цикла в виде строковых обозначений, мы получим einsum решение на основе, как так -

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

Будучи в основном вопросом, связанным с расширением (поскольку мы не уменьшаем ни одной оси), мы можем просто использовать NumPy broadcasting тоже, вводя новые оси с None/np.newaxis в разных местах для расширения, например,

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