Как эффективно вычислить внешний продукт двух серий матриц в numpy?

Допустим, что у меня есть матрицы A (KxMxN) и B (KxLxN), где L, M, N малы, а K – большое число. Я хотел бы рассчитать внешний продукт использования последних двух измерений вдоль первого измерения, чтобы получить матрицу C (KxMxL).

Я могу сделать это, запустив цикл for для каждого индекса k в «K» и используя функцию matmul numpy для двумерных матриц

out = [np.matmul(x,yT) for x, y in zip(A, B)] out=np.asarray(out) 

Интересно, могу ли я сделать это без цикла / понимания, поскольку K – очень большое число.

Поскольку A имеет форму (K, M, N) и B имеет форму (K, L, N) , и вы хотите найти сумму продуктов с формой (K, M, L) , вы можете использовать np.einsum :

 C = np.einsum('kmn,kln->kml', A, B) 

matmul работает с транспозицией в B , поэтому его второй до последнего тусклый совпадает с последним из A

 In [1019]: A=np.random.rand(K,M,N) In [1021]: B=np.random.rand(K,L,N) In [1023]: C=np.einsum('kmn,kln->kml',A,B) In [1024]: C.shape Out[1024]: (2, 4, 3) In [1026]: D=A@B.transpose(0,2,1) In [1027]: D.shape Out[1027]: (2, 4, 3) In [1028]: np.allclose(C,D) Out[1028]: True 

Для этого небольшого примера timeit и тот же.

[np.dot(x,yT) for x, y in zip(A, B)] делает то же самое; совпадают с 2-м последним тусклым y , равным 2, и итерацией 1-го разряда A и B