einsum

Одно правило

Нотация Эйнштейна выглядит как набор рецептов, а является одним правилом:

$$\htmlData{k=rule}{c_{ik} = \sum_j a_{ij} b_{jk}} \qquad\longleftrightarrow\qquad \texttt{'ij,jk->ik'}$$

Правило целиком: индекс, который есть во входах и отсутствует в выходе, суммируется. Всё остальное — следствия.

Отсюда весь «список рецептов» перестаёт быть списком:

спецификация	что получается	почему
'ij,jk->ik'	матричное умножение	j пропал → сумма по j
'ij,j->i'	матрица на вектор	то же
'i,i->'	скалярное произведение	i пропал, выход скаляр
'i,j->ij'	внешнее произведение	ничего не пропало → суммы нет
'ij->ji'	транспонирование	ничего не пропало, только порядок
'ij->i'	суммы по строкам	j пропал
'ij->'	сумма всего	пропали оба
'ii->i'	диагональ	повтор индекса = выбор диагонали
'ii->'	след	диагональ, затем сумма

Последние две строки показывают вторую половину нотации: повторённый индекс внутри одного операнда означает диагональ. 'ii->' — это $\sum_i a_{ii}$, то есть след.

Цена операции читается глазами

Практическая ценность нотации не в краткости, а в том, что по ней сразу виден объём работы.

Число умножений-сложений равно произведению размеров всех индексов — и суммируемых, и остающихся. Причина простая: одно умножение-сложение происходит на каждый набор значений всех индексов.

Проверьте на виджете: у 'ij,jk->ik' при $i=2, j=3, k=4$ выходит $24$ операции, и это $2 \cdot 3 \cdot 4$. Результат содержит $8$ элементов, значит на каждый приходится по $3$ операции — ровно длина суммируемой оси.

ij,jk->ikij->jiii->bij,bjk->bikbhqd,bhkd->bhqkbhqk,bhkd->bhqd

ij, jk → ik

индекс	размер	роль
i	2	остаётся
j	3	суммируется
k	4	остаётся

форма результата

[2, 4]

элементов

8

умножений-сложений

24

операций на элемент

3

Столбец «операций на элемент» полезен сам по себе. Он равен произведению суммируемых размеров, и по нему видно, чем ограничена операция:

• отношение $1$ — суммирования нет, операция только двигает данные. Транспонирование, перестановка осей, поэлементные операции. Ограничены пропускной способностью памяти;
• большое отношение — много арифметики на каждое прочитанное число. Матричные умножения. Ограничены вычислительной мощностью.

Это и есть различие между memory-bound и compute-bound, увиденное из нотации.

Почему порядок свёртки важен

Возьмём произведение трёх матриц $A_{i j} B_{j k} C_{k l}$ с размерами $i = 1000$, $j = 10$, $k = 10$, $l = 1000$.

порядок	промежуточный результат	операций
$(AB)C$	$i \times k = 10^4$	$i j k + i k l = 10^5 + 10^7$
$A(BC)$	$j \times l = 10^4$	$j k l + i j l = 10^5 + 10^7$

Здесь одинаково. А теперь $i = 1000$, $j = 1000$, $k = 2$, $l = 1000$:

порядок	операций
$(AB)C$	$i j k + i k l = 2\cdot10^6 + 2\cdot10^6 = 4\cdot10^6$
$A(BC)$	$j k l + i j l = 2\cdot10^6 + 10^9$

Разница в $250$ раз. Причина — узкая ось $k$: свернув по ней первым делом, вы избегаете создания матрицы $j \times l$.

Отсюда важное предупреждение: einsum с тремя и более операндами не гарантирует оптимального порядка. NumPy принимает optimize=True, opt_einsum подбирает порядок поиском; наивная реализация может выбрать худший вариант. Для критичных мест порядок стоит задавать явно, разбив выражение на два вызова.

Что einsum не делает

Честные ограничения:

• не оптимизирует за вас порядок свёртки при трёх и более операндах, если не попросить;
• не всегда быстрее специализированных вызовов: matmul может использовать более подходящее ядро, чем то, во что развернётся einsum;
• не поддерживает всё, что нужно: einops появился потому, что einsum не умеет разбивать и склеивать оси ('b (h d) -> b h d').

Но у него есть то, чего нет ни у одной специализированной функции: спецификация читается как утверждение о формах. 'bhqd,bhkd->bhqk' говорит, что головы и батч сохраняются, а свёртка идёт по размеру головы, — и если вы ошиблись осью, ошибка видна в строке, а не в отладчике.