Удивительное дело, но в русскоязычном сегменте интернета почти нет материала, разъясняющего понятным языком соглашение Эйнштейна о суммировании. Не менее удивительно то, что материалов, позволяющих понять принцип работы функции einsum в русскоязычном интернете ещё меньше. На английском есть довольно развёрнутый ответ о работе einsum на stack overflow, а на русском только некоторое число сайтов, предоставляющих кривой перевод этого самого ответа. Хочу исправить эту проблему с недостатком материалов, и всех, кому интересно приглашаю к прочтению!
Обсуждаем соглашение Эйнштейна
Прежде всего отмечу, что соглашение Эйнштейна чаще всего используется в тензорном анализе и его приложениях, поэтому дальше в статье будет несколько референсов к тензорам.
Когда вы только начинаете работать с тензорами, вас может смутить, что кроме привычных подстрочных индексов, используются также и надстрочные индексы, которые по началу вообще можно принять за возведение в степень. Пример:
"а с верхним индексом i" будет записано как , а "a в квадрате с верхним индексом i" будет записываться . Возможно, по-началу это вводит в заблуждение и кажется неудобным, но со временем можно привыкнуть.
Соглашение: далее в статье объекты вида или я буду называть термами.
О чём вообще соглашение Эйнштейна?
Соглашение Эйнштейна призвано уменьшить число знаков суммирования в выражении. Есть три простых правила, определяющие, насколько то или иное выражение корректно записано в нотации Эйнштейна.
Правило № 1: Суммирование ведётся по всем индексам, повторяющимся дважды в одном терме.
Пример: рассмотрим выражение следующего вида:
С использованием соглашения Эйнштейна это выражение может быть переписано так:
Таким образом мы избавляемся от знака суммы, и просто пишем единственный терм. Обратим внимание, что в этом терме индекс i повторяется дважды, а значит, в соответствие с первым правилом мы понимаем, что суммирование ведётся по индексу i, а точнее, по всем возможным значениям, которые принимает этот индекс.
Рассмотрим ещё один пример: пусть нам нужно умножить матрицу на вектор . Результатом будет являться вектор . По определению:
Соглашение Эйнштейна позволяет избавиться от знака суммы:
Заметим, что в терм индекс i входит один раз, а индекс j входит два раза, а значит, суммирование ведётся по индексу j.
Определение 1. Индекс, который входит в терм дважды, называется фиктивным индексом.
Определение 2. Свободным индексом назовём все индексы в терме, не являющие фиктивными.
Отметим, что каждый фиктивный индекс может быть заменён любым другим фиктивным индексом, при условии, что
-
Новый фиктивный индекс не входит в множество свободных индексов терма.
-
Новый фиктивный индекс принимает то же множество значений, что и старый фиктивный индекс.
Чтобы объяснить проще, рассмотрим следующий код на языке Python:
for i in range(M):
for j in range(N):
b[i, j] = A[i, j] * v[j]
Этот код кратко описывает процесс умножения матрицы на вектор, а точнее, этот пример. Здесь индекс j является фиктивным, а индекс i – свободным. Суммирование в соглашении Эйнштейна ведётся по фиктивным индексам. Имя переменной j мы можем заменить на любое другое.
Правило № 2. В каждом терме не может встречаться более двух одинаковых индексов.
Второе правило говорит нам, что мы можем написать , но не можем написать или , несмотря на то, что на практике такие выражения всё же имеют смысл.
Больше примеров:
– здесь i является фиктивным индексом, т.к. повторяется дважды;
– здесь i является свободным индексом, а j – фиктивным;
– здесь и i, и j являются фиктивными индексами;
– здесь и i, и j являются фиктивными индексами;
– не правильно по второму правилу (индекс i входит в терм трижды);
Из примеров выше можно заключить, что когда мы считаем число вхождений индексов в терм, мы не делаем разницы между верхними и нижними индексами, и считаем их вместе. Ещё один важный пример: когда мы видим выражение следующего вида
Мы должны понимать, что это выражение записано верно, и не противоречит второму правилу. Действительно, если посчитать все вхождения индексов, то получится, что индекс i входит 3 раза, как и индекс j, но в выражении записано два терма, а не один, и если посчитать вхождение индексов в каждый терм отдельно (как того и требует второе правило), то мы увидим, что ничего не нарушается.
Правило № 3. В уравнениях, записанных с использованием соглашения Эйнштейна свободные индексы слева и свободные индексы справа должны совпадать.
Рассмотрим несколько примеров для закрепления этого правила:
– этот пример мы уже рассматривали выше, здесь i является свободным индексом левой части уравнения, и свободным индексом правой части уравнения;
– пример посложнее. Посчитаем вхождения индексов для каждого терма: в первый терм правой части k и j входят дважды, значит, они являются фиктивными индексами, i входит один раз, значит, является свободным. Во второй терм правой части k входит два раза, i – один, значит, k – фиктивный, i – свободный. В левой части индекс i входит один раз, а значит, является свободным. Итог: индекс i является свободным для обеих частей уравнения, а значит, правило 3 выполнено.
Рассмотрим так же несколько примеров, в которых третье правило не выполняется:
– слева i является свободным индексом, но справа свободны индексы i и j;
– слева свободен индекс j, но справа свободен индекс i. Свободные индексы не совпадают;
– здесь слева свободен индекс i, а справа свободны индексы i, j;
Пример упрощения сложного выражения с помощью соглашения Эйнштейна: тензорный поезд
Пусть – пятимерный тензор. Тогда утверждается, что он может быть представлен в следующем виде:
Там сейчас не очень важно, что из себя представляется каждая , и что такое . Наша задача сейчас – исключительно синтаксическая игра. Нужно упростить выражение, особо не вникая в смысл происходящего.
Прежде всего видно, что свободными индексами являются , а фиктивными, соответственно индексы . Расположим индексы в соседних множителях так, чтобы в первом множителе индекс, по которому идёт суммирование, стоял снизу, а во втором тот же самый индекс стоял сверху. Так же заметим, что множителями являются тензоры , и у них в верхнем регистре уже стоит . Чтобы повысить читаемость, будем оборачивать множители в скобки, и только потом ставить индексы. Само же упрощённое выражение переписывается из исходного почти дословно:
Ура, мы научились упрощать сложные выражения с помощью соглашения Эйнштейна!
Обсуждаем einsum
einsum это функция, присутствующая в нескольких популярных библиотеках для Python (NumPy, TensorFlow, PyTorch). Во всех библиотеках, в которых эта функция реализована, она работает одинаково (с точностью до функционала структур, определённых в конкретной библиотеке), поэтому нет смысла рассматривать один и тот же пример в разных библиотеках, достаточно рассказать про einsum в одной конкретной библиотеке. Далее в статье я буду использовать NumPy. einsum применяет соглашение Эйнштейна о суммировании к переданным массивам. Функция принимает множество опциональных аргументов, про них лучше почитать в документации, мы же сейчас разберём, как передавать шаблон, по которому функция будет применять соглашение Эйнштейна.
Рассмотрим сразу такой пример: пусть , – две матрицы, и мы хотим их перемножить. Результатом будет матрица , которую мы можем записать следующим образом, используя определение матричного умножения и соглашение Эйнштейна:
Теперь пусть мы хотим перемножить их программно. Ну, это можно довольно просто сделать с помощью трёх вложенных циклов:
M = np.zeros((3, 2))
for i in range(3):
for j in range(2):
for k in range(5):
M[i, j] = A[i, k] * B[k, j]
Либо, используя функцию einsum можно написать это произведение в одну строчку:
M = np.einsum("ik,kj->ij", A, B)
Разберёмся, что за магия происходит в этой строчке. einsum принимает один обязательный аргумент: шаблон, по которому будет применено соглашение Эйнштейна. Шаблон этот выглядит так:
"{индексы, определяющие размерность первого массива},{индексы, определяющие размерность второго массива}->{индексы, определяющие размерность результирующего массива}"
Поведение шаблона einsum определяется следующими правилами:
-
Если один и тот же индекс встречается слева и справа от запятой (до стрелочки), то суммирование будет вестись по этому индексу;
-
Если после стрелочки ничего не написано, то суммирование произойдёт по всем встреченным осям;
-
Никакой индекс не должен встречаться 3 и более раз;
Таким образом мы видим, что einsum очень естественно поддерживает понятие свободных и фиктивных индексов, а также первые два правила, которые мы вводили, пока обсуждали соглашение Эйнштейна. Кроме того, как выражение, написанное с помощью соглашения Эйнштейна, может быть развёрнуто с помощью введения знаков суммы, так и функция einsum может быть развёрнута с помощью нескольких вложенных циклов. Это может быть очень удобно на первых порах, пока не сформируется устойчивое понимание einsum.
Рассмотрим теперь некоторое количество примеров разной степени сложности, чтобы закрепить понимание einsum:
Одна einsum, чтобы править всеми
Пример 1. Сумма всех значений вектора:
vector = np.array([1, 2, 3, 4, 5])
result = np.einsum("i->", vector)
print(result)
Output
15
Пример 2. Сумма всех значений матрицы:
matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->", matrix)
print(result)
Output
21
Пример 3. Сумма значений по столбцам:
matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->j", matrix)
print(result)
Output
[9, 12]
Пример 4. Сумма значений по строкам:
matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->i", matrix)
print(result)
Output
[3, 7, 11]
Пример 5. Транспонирование (я об этом не написал, но оси, по которым суммирования не произошло, мы можем возвращать в любом порядке):
matrix = np.array([[1, 2], [3, 4], [5, 6]])
result = np.einsum("ij->ji", matrix)
print(result)
Output
[[1, 3, 5], [2, 4, 6]]
Пример 6. Умножение матрицы на вектор:
matrix = np.array([[1, 2], [3, 4], [5, 6]])
vector = np.array([[1, 2]])
result = np.einsum("ij,kj->ik", matrix, vector)
print(result)
Заметим, что вектор имеет форму , и чтобы умножить матрицу на него по правилам, его нужно было бы транспонировать. Однако с помощью einsum мы можем задать ось, по которой будет вестись суммирование, и немного выиграть по памяти, не создавая копию уже существующего вектора.
Output
[[5], [11], [17]]
Пример 7. Умножение матрицы на матрицу:
matrix1 = np.array([[1, 2], [3, 4], [5, 6]])
matrix2 = np.array([[1, 0], [0, 1]])
result = np.einsum("ik,kj->ij", matrix1, matrix2)
print(result)
Output
[[1, 2], [3, 4], [5, 6]]
Пример 8. Скалярное произведение векторов:
vector1 = np.array([[1, 2, 3]])
vector2 = np.array([[1, 1, 1]])
result = np.einsum("ik,jk->", vector1, vector2)
print(result)
Output
6
Пример 9. След матрицы:
matrix1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = np.einsum("ii->", matrix1)
print(result)
Output
15
Пример 10. Адамарово (покомпонентное) произведение:
matrix1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
matrix2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
result = np.einsum("ij,ij->ij", matrix1, matrix2)
print(result)
Это может показаться контринтутивно, но, как написано выше: если не понятно, что делает einsum – запиши через циклы:
result = np.zeros(matrix1.shape, dtype="int32")
for i in range(result.shape[0]):
for j in range(result.shape[1]):
result[i, j] += matrix1[i, j] * matrix2[i, j]
print(result)
Output
[[1, 0, 0], [0, 5, 0], [0, 0, 9]]
Пример 11. Кронекерово (внешнее) произведение векторов:
vector1 = np.array([1, 2, 3])
vector2 = np.array([1, 0, 0])
result = np.einsum("i,j->ij", vector1, vector2)
print(result)
Output
[[1, 0, 0], [2, 0, 0], [3, 0, 0]]
Пример 12. Транспонирование тензора:
A = np.array([[[0, 1], [1, 2], [2, 3]], [[1, 2], [2, 3], [3, 4]], [[2, 3], [3, 4], [4, 5]]])
result = np.einsum("ijk->jki", A)
print(result)
Output
[[[0, 1, 2], [1, 2, 3]], [[1, 2, 3], [2, 3, 4]], [[2, 3, 4], [3, 4, 5]]]
Пример 13. Произведение тензора на матрицу по третьей моде:
A = np.array([[[0, 1], [1, 2], [2, 3]], [[1, 2], [2, 3], [3, 4]], [[2, 3], [3, 4], [4, 5]]])
U = np.array([[1, 2], [2, 3]])
result = np.einsum("ijk,nk->ijn", A, U)
print(result)
Output
[[[2, 3], [5, 8], [8, 13]], [[5, 8], [8, 13], [11. 18]], [[8, 13], [11, 18], [14, 23]]]
Итоги
Конечно, einsum поставляет только дополнительный синтаксический сахар. Всегда можно использовать цепочки вложенных циклов, множество библиотечных функций (np.dot, np.outer, np.tensordot, np.transpose, np.cumsum и т.д.), и вообще не использовать einsum. Но если потратить время и понять, как она работает, то можно научиться писать гораздо более сжатый, и, не побоюсь этого слова, эффективный код.
Ссылки
Ролик с примерами einsum (ещё больше примеров).
Соглашение Эйнштейна (база)
Соглашение Эйнштейна (продвинутая часть)
Автор: Ваня