Суммирование Эйнштейна работает просто отлично:
>>> a = np.random.randn(100, 4, 4)
>>> b = np.random.randn(100, 4)
>>> foo = np.einsum('ijk,ik->ij', a, b)
>>> bar = np.zeros_like(foo)
>>> for i, (ai, bi) in enumerate(zip(a, b)):
bar[i] = np.dot(ai, bi)
>>> np.allclose(foo, bar)
True
Чтобы немного пояснить суммирование, обратите внимание, что вы сжимаете last ось b
. Таким образом, вы можете думать о выполнении каждого внутреннего продукта, как если бы np.einsum('jk,k->j', a[0], b[0])
Но мы делаем по одному для каждого элемента a
и b
, таким образом, включается первая ось, которая не сокращена. Следовательно, ijk,ik->ij
.