In [601]: C
Out[601]: array([436., 534., 644., 766., 900.])
Это естественно для einsum
.Я использую i
, как вы, чтобы обозначить индекс, который влияет на результат.j
и k
- это индексы, которые используются в сумме произведений.
In [602]: np.einsum('ji,jk,ki->i',A,B,A)
Out[602]: array([436, 534, 644, 766, 900])
Возможно, это также можно сделать с помощью mutmul
, хотя для этого может потребоваться добавление размера и последующее сжатие.
dot
подходы, использующие diag
, выполняют намного больше работы, чем необходимо.diag
выбрасывает много значений.
Чтобы использовать matmul
, мы должны сделать измерение i
первым из трехмерных массивов.Это «пассивный» результат, который переносится на результат:
In [603]: A.T[:,None,:]@B@A.T[:,:,None]
Out[603]:
array([[[436]], # (5,1,1) result
[[534]],
[[644]],
[[766]],
[[900]]])
In [604]: (A.T[:,None,:]@B@A.T[:,:,None]).squeeze()
Out[604]: array([436, 534, 644, 766, 900])
Или индексировать дополнительные измерения: (A.T[:,None,:]@B@A.T[:,:,None])[:,0,0]