Эффективно вычислить скалярное произведение на последнем измерении массива

Question

Какой самый быстрый способ вычисления точечного произведения на последнем измерении многомерного ndarray?

На данный момент я делаю это:

import numpy as np

a=np.reshape(np.arange(90),[3,3,2,5])
b=np.reshape(np.arange(90),[3,3,2,5])
# for the sake of simplicity, a and b are the same for this example

ab=(a*b).sum(axis=-1)

Я думаю, что einsum может быть полезно здесь, но я изо всех сил пытаюсь применить это к моему случаю.

Спасибо!

Answer 1

Для универсальных массивов ndim для получения суммовых сокращений по последней оси -

np.einsum('...i,...i->...',a,b)

Альтернатива с np.matmul -

np.matmul(a[...,None,:],b[...,None])[...,0,0]

Примечание: На Python 3.x np.matmulможно заменить на @ operator.