Python: векторизация через массив numpy

Question

X и Y - это трехмерные массивы с размерами (a,b,c). Моя цель - сделать точечный продукт.

Рассмотрим случай, когда индексы i и j являются скалярными, а (X[i,:,j].T).dot(Y[i,:,j]) будет простым и вернет скаляр.

Однако, если я попытаюсь сделать векторизацию, i и j станут 1d-массивами, а (X[i,:,j].T).dot(Y[i,:,j]) вернут матрицу, но я ожидаю в качестве результата 1d-массив. Как мне обойти эту проблему?

Answer 1

Наивная реализация с использованием понимания списка:

a,b,c = X.shape

r1 = [(X[i,:,j].T).dot(Y[i,:,j]) for i in range(a) for j in range(c)]

Реализация с использованием np.einsum :

r2 = np.einsum('ijk,ijk->ik', X,Y).flatten()