einsum — NumPy — Тренажёр

Коротко

np.einsum("ij,jk->ik", A, B) описывает вычисление буквами: у каждого операнда свои индексы, после -> — индексы результата.
Индекс, который есть во входах, но нет в результате, суммируется. Индекс, повторённый в разных операндах, перемножается.
"ij->ji" — транспонирование, "ii->" — след, "ij->j" — сумма по строкам, "i,i->" — скалярное произведение.
Удобно для пачек и многомерных тензоров, где @ и sum дают громоздкий код.
optimize=True подбирает порядок вычислений для трёх и более операндов.

Синтаксис

np.einsum("индексы_a,индексы_b->индексы_результата", a, b, optimize=False)
np.einsum("ij,jk->ik", A, B)      # A @ B
np.einsum("bij,bjk->bik", X, Y)   # пачка матриц
np.einsum("ij->j", A)             # A.sum(axis=0)

Запись	Эквивалент
`"ij->ji"`	`A.T`
`"ii->"`	`np.trace(A)`
`"ij->"`	`A.sum()`
`"ij->j"`	`A.sum(axis=0)`
`"ij,jk->ik"`	`A @ B`
`"i,i->"`	`v @ v`
`"i,j->ij"`	`np.outer(v, w)`
`"ij,ij->ij"`	`A * B`

Примеры

Данные для примеровиспользуются во всех примерах ниже

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])
v = np.array([1, 10])

Суммы

print(np.einsum("ij->", A))        # сумма всех элементов
print(np.einsum("ij->j", A))       # сумма по строкам: для каждого столбца
np.einsum("ij->i", A)              # для каждой строки

Вывод

10
[4 6]
array([3, 7])

Транспонирование и след

print(np.einsum("ij->ji", A))      # транспонирование
np.einsum("ii->", A)               # след: сумма диагонали

Вывод

[[1 3]
 [2 4]]
np.int64(5)

Произведения

print(np.einsum("ij,jk->ik", A, B))   # матричное произведение, как A @ B
print(np.einsum("ij,j->i", A, v))     # матрица на вектор
np.einsum("i,i->", v, v)              # скалярное произведение

Вывод

[[19 22]
 [43 50]]
[21 43]
np.int64(101)

Поэлементно и построчно

print(np.einsum("ij,ij->ij", A, B))   # поэлементно, как A * B
np.einsum("ij,ij->i", A, B)           # построчное скалярное произведение

Вывод

[[ 5 12]
 [21 32]]
array([17, 53])

Пачка матриц

batch_a = np.ones((5, 2, 3))
batch_b = np.ones((5, 3, 4))
np.einsum("bij,bjk->bik", batch_a, batch_b).shape   # пачка матриц

Вывод

(5, 2, 4)

Внешнее произведение

np.einsum("i,j->ij", v, v)            # внешнее произведение

Вывод

array([[  1,  10],
       [ 10, 100]])

Подводные камни

Одна буква — один размер

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

Разные размеры у одного индексаошибка: ValueError

np.einsum("ij,jk->ik", A, np.ones((3, 2)))   # j = 2 и j = 3 не совпадают

Ошибка

ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (2,2)->(2,newaxis,2) (3,2)->(2,3)

Читаемость

Для простых операций @, .T и .sum(axis=…) понятнее. einsum оправдан, когда операция не выражается одной стандартной функцией.