Как умножить два 3D-массива numpy вдоль некоторой определенной оси?

#python #arrays #numpy

#python #массивы #numpy

Предположим, у нас есть два массива numpy: A с формой (n,p,q) , B с формой (n,q,r) .

Как их умножить, чтобы получить массив C с формой (n,p,r) ? Я имею в виду сохранить ось 0 и умножить их на оси 1 и 2.

Я знаю, что это можно вычислить с помощью: C = np.stack([np.dot(a[i], b[i]) for i in range(A.shape[0])])

Но существует ли функция numpy, которую можно использовать для ее непосредственного вычисления?

Я думаю, вы можете сделать np.einsum :

 # sample data
n,p,q,r = 2,3,4,5
A = np.zeros((n,p,q))
B = np.zeros((n,p,r))

out = np.einsum('npq,nqr->npr',A,B)

out.shape
# (2, 3, 5)