torch.einsum

torch.einsum(equation, *operands) → Tensor [source]

Вычисляет сумму произведения элементов входных operands по измерениям, указанным с помощью обозначений, основанных на соглашении об обозначении Эйнштейна.

Einsum позволяет вычислять многие распространённые многомерные линейно-алгебраические операции над массивами, представляя их в сокращённой форме на основе соглашения об обозначении Эйнштейна, заданного equation. Подробности этого формата описаны ниже, но основная идея заключается в том, чтобы пометить каждое измерение входных operands определённым индексом и указать, какие индексы являются частью результата. Затем результат вычисляется путём суммирования произведения элементов operands по измерениям, индексы которых не являются частью результата. Например, произведение матриц может быть вычислено с помощью einsum как torch.einsum(“ij,jk->ik”, A, B). Здесь j является индексом суммирования, а i и k — индексами результата (подробнее о том, почему, см. ниже).

Уравнение:

Строка equation задаёт индексы (буквы в [a-zA-Z]) для каждого измерения входных operands в том же порядке, что и измерения, разделяя индексы для каждого операнда запятой (','), например, ‘ij,jk’ задают индексы для двух 2D операндов. Измерения с одинаковым индексом должны быть согласованы по размеру, то есть их размер должен совпадать или быть 1. Исключением является случай, когда индекс повторяется для одного и того же входного операнда, в этом случае измерения, помеченные этим индексом для этого операнда, должны совпадать по размеру, и операнд будет заменён своей диагональю по этим измерениям. Индексы, которые встречаются ровно один раз в equation , будут частью результата, отсортированного в порядке возрастания по алфавиту. Результат вычисляется путём умножения входных operands поэлементно, с выравниванием измерений на основе индексов, а затем суммирования измерений, индексы которых не являются частью результата.

Необязательно, индексы результата могут быть явно определены путём добавления стрелки (’->’) в конце уравнения, за которой следуют индексы результата. Например, следующее уравнение вычисляет транспонированное произведение матриц: ‘ij,jk->ki’. Индексы результата должны встречаться как минимум один раз для некоторого входного операнда и не более одного раза для результата.

Эллипсис (’…’) может быть использован вместо индексов для согласования измерений, охваченных эллипсом. Каждый входной операнд может содержать не более одного эллипса, который будет охватывать измерения, не охваченные индексами, например, для входного операнда с 5 измерениями эллипсис в уравнении ‘ab…c’ охватывает третье и четвёртое измерения. Эллипсис не обязательно должен охватывать одинаковое количество измерений по всем operands, но «форма» эллипса (размер измерений, охваченных ими) должна быть согласована. Если результат не определён явно с помощью обозначения со стрелкой (’->’), эллипсис будет стоять первым в результате (самые левые измерения), перед индексами, которые встречаются ровно один раз для входных операндов. Например, следующее уравнение реализует пакетное умножение матриц ‘…ij,…jk’.

Несколько заключительных замечаний: в уравнении могут содержаться пробелы между различными элементами (индексами, эллипсами, стрелкой и запятыми), но что-то вроде ‘…’ не является допустимым. Пустая строка ‘’ допустима для скалярных операндов.

Примечание

torch.einsum обрабатывает эллипсы (’…’) по-другому, чем NumPy, позволяя измерениям, охватываемым эллипсом, суммироваться, то есть эллипсис не обязателен для результата.

Примечание

Данная функция использует opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/) для ускорения вычислений или экономии памяти за счёт оптимизации порядка сокращения. Эта оптимизация происходит, когда входных данных три или более, так как порядок в противном случае не важен. Обратите внимание, что поиск _оптимального_ пути является задачей NP-сложности, поэтому opt_einsum опирается на различные эвристики для достижения почти оптимальных результатов. Если opt_einsum недоступен, по умолчанию используется порядок сокращения слева направо.

Чтобы обойти это поведение по умолчанию, добавьте следующую строку, чтобы отключить использование opt_einsum и пропустить вычисление пути: torch.backends.opt_einsum.enabled = False

Чтобы указать стратегию, которую вы хотите использовать для opt_einsum при вычислении пути сокращения, добавьте следующую строку: torch.backends.opt_einsum.strategy = ‘auto’. По умолчанию используется стратегия 'auto', и мы также поддерживаем 'greedy' и 'optimal'. Отметим, что время работы 'optimal' факториально зависит от числа входных данных! Дополнительные сведения см. в документации opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/path_finding.html).

Примечание

Начиная с PyTorch 1.10 torch.einsum() также поддерживает формат подсписков (см. примеры ниже). В этом формате индексы для каждого операнда задаются подсписками, списками целых чисел в диапазоне [0, 52). Эти подсписки следуют за своими операндами, и дополнительный подсписок может появиться в конце входных данных для указания индексов результата, например torch.einsum(op1, sublist1, op2, sublist2, …, [subslist_out]). Объект Python Ellipsis может быть предоставлен в подсписке для включения согласования по размеру, как описано в разделе «Уравнение» выше.

Параметры:

• equation (str) – Индексы для суммирования Эйнштейна.
• operands (List[Tensor]) – Тензоры для вычисления суммирования Эйнштейна.

Тип возвращаемого значения:

Tensor

Примеры:

>>> # trace
>>> torch.einsum('ii', torch.randn(4, 4))
tensor(-1.2104)

>>> # diagonal
>>> torch.einsum('ii->i', torch.randn(4, 4))
tensor([-0.1034,  0.7952, -0.2433,  0.4545])

>>> # outer product
>>> x = torch.randn(5)
>>> y = torch.randn(4)
>>> torch.einsum('i,j->ij', x, y)
tensor([[ 0.1156, -0.2897, -0.3918,  0.4963],
        [-0.3744,  0.9381,  1.2685, -1.6070],
        [ 0.7208, -1.8058, -2.4419,  3.0936],
        [ 0.1713, -0.4291, -0.5802,  0.7350],
        [ 0.5704, -1.4290, -1.9323,  2.4480]])

>>> # batch matrix multiplication
>>> As = torch.randn(3,2,5)
>>> Bs = torch.randn(3,5,4)
>>> torch.einsum('bij,bjk->bik', As, Bs)
tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],
        [-1.6706, -0.8097, -0.8025, -2.1183]],

        [[ 4.2239,  0.3107, -0.5756, -0.2354],
        [-1.4558, -0.3460,  1.5087, -0.8530]],

        [[ 2.8153,  1.8787, -4.3839, -1.2112],
        [ 0.3728, -2.1131,  0.0921,  0.8305]]])

>>> # with sublist format and ellipsis
>>> torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2])
tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],
        [-1.6706, -0.8097, -0.8025, -2.1183]],

        [[ 4.2239,  0.3107, -0.5756, -0.2354],
        [-1.4558, -0.3460,  1.5087, -0.8530]],

        [[ 2.8153,  1.8787, -4.3839, -1.2112],
        [ 0.3728, -2.1131,  0.0921,  0.8305]]])

>>> # batch permute
>>> A = torch.randn(2, 3, 4, 5)
>>> torch.einsum('...ij->...ji', A).shape
torch.Size([2, 3, 5, 4])

>>> # equivalent to torch.nn.functional.bilinear
>>> A = torch.randn(3,5,4)
>>> l = torch.randn(2,5)
>>> r = torch.randn(2,4)
>>> torch.einsum('bn,anm,bm->ba', l, A, r)
tensor([[-0.3430, -5.2405,  0.4494],
        [ 0.3311,  5.5201, -3.0356]])