Функция einsum_path

Функция einsum_path анализирует выражение einsum и возвращает оптимальный порядок выполнения операций с тензорами. Первым параметром передаётся строка с подстрочной нотацией, далее передаются входные массивы. Функция возвращает кортеж из двух элементов: список путей и строку с подробной информацией о вычислениях.

Синтаксис

np.einsum_path(subscripts, *operands, optimize='greedy')

Параметр optimize может принимать значения 'greedy' (жадный алгоритм) или 'optimal' (полный перебор). По умолчанию используется 'greedy', который даёт хороший баланс между скоростью и качеством оптимизации.

Пример

Создадим три матрицы и найдём оптимальный порядок их перемножения:

import numpy as np

rng = np.random.default_rng(0)
arr1 = rng.random((2, 3))
arr2 = rng.random((3, 4))
arr3 = rng.random((4, 5))

path_info = np.einsum_path('ij,jk,kl->il', arr1, arr2, arr3)
print(path_info[0])

Результат выполнения кода:

['einsum_path', (0, 1), (0, 1)]

Пример

Выведем подробную информацию о ходе оптимизации с помощью флага optimize:

import numpy as np

rng = np.random.default_rng(0)
arr1 = rng.random((3, 4))
arr2 = rng.random((4, 5))
arr3 = rng.random((5, 6))

path, desc = np.einsum_path('ij,jk,kl->il', arr1, arr2, arr3)
print(desc)

Результат выполнения кода:

"  Complete contraction:  ij,jk,kl->il
		 Naive scaling:  4
	 Optimized scaling:  3
	  Naive FLOP count:  1.200e+03
  Optimized FLOP count:  4.800e+02
   Theoretical speedup:  2.500
  Largest intermediate:  2.400e+01 elements
---------------------------------------------------------------------------
									  current							   
								 ~~~~~~~~~~~~~~~						   
  operation					  size	memory	  naive	   optimized
								  ^		 ^		 ^			 ^
  'ij,jk->ik'				 (3,4)		1   1.200e+03	 4.800e+02
  'ik,kl->il'				 (3,5)		2   4.800e+02	 4.800e+02"

Пример

Применим полученный путь для выполнения einsum с оптимизацией:

import numpy as np

rng = np.random.default_rng(0)
arr1 = rng.random((2, 3))
arr2 = rng.random((3, 4))
arr3 = rng.random((4, 5))

path, _ = np.einsum_path('ij,jk,kl->il', arr1, arr2, arr3)
res = np.einsum('ij,jk,kl->il', arr1, arr2, arr3, optimize=path)
print(res.shape)

Результат выполнения кода:

(2, 5)

Пример

Сравним производительность с жадным и полным алгоритмом оптимизации для четырёх матриц:

import numpy as np

rng = np.random.default_rng(0)
arr1 = rng.random((2, 3))
arr2 = rng.random((3, 4))
arr3 = rng.random((4, 5))
arr4 = rng.random((5, 6))

path_greedy = np.einsum_path('ij,jk,kl,lm->im', arr1, arr2, arr3, arr4, optimize='greedy')[0]
path_optimal = np.einsum_path('ij,jk,kl,lm->im', arr1, arr2, arr3, arr4, optimize='optimal')[0]

print(path_greedy)
print(path_optimal)

Результат выполнения кода:

['einsum_path', (0, 1), (0, 1), (0, 1)]
['einsum_path', (0, 1), (0, 1), (0, 1)]

Смотрите также

функцию einsum,
которая выполняет тензорные сокращения по подстрочной нотации
функцию tensordot,
которая вычисляет тензорное произведение с указанием осей
функцию multi_dot,
которая вычисляет произведение нескольких матриц с оптимизацией порядка
функцию dot,
которая вычисляет скалярное произведение двух массивов