Multiplicação de matrizes com dependência do iterador - NumPy

Algum tempo atrásthis question (agora excluído, mas os usuários com mais de 10 mil representantes ainda podem visualizá-lo) foi postado. Pareceu-me interessante e aprendi algo novo lá ao tentar resolvê-lo e achei que valeria a pena compartilhar. Gostaria de postar essas idéias / soluções e gostaria de ver as pessoas postando outras maneiras possíveis de resolvê-lo. Estou postando a essência da pergunta a seguir.

Então, nós temos dois ndarrays NumPya eb de formas:

a : (m,n,N)
b : (n,m,N)

Vamos supor que estamos lidando com casos em quem,n & N são comparáveis.

O problema é resolver a seguinte multiplicação e soma com foco no desempenho:

def all_loopy(a,b):
    P,Q,N = a.shape
    d = np.zeros(N)
    for i in range(N):
        for j in range(i):
            for k in range(P):
                for n in range(Q):
                    d[i] += a[k,n,i] * b[n,k,j]
    return d

questionAnswers(2)

yourAnswerToTheQuestion