在 Python 中评估 einsum 表达式的最低成本收缩顺序

pythonnumpyserver side programmingprogramming更新于 2024/5/14 16:44:17

要获取 einsum 表达式的最低成本收缩顺序,请使用 Python 中的 numpy.einsum+path() 方法。第一个参数下标指定求和的下标。第二个参数操作数是操作的数组。

使用爱因斯坦求和约定,许多常见的多维线性代数数组运算可以以简单的方式表示。在隐式模式下,einsum 计算这些值。

在显式模式下,einsum 通过禁用或强制对指定下标标签求和,提供了进一步的灵活性来计算可能不被视为经典爱因斯坦求和运算的其他数组运算。

结果路径指示应首先收缩输入收缩的哪些项,然后将此收缩的结果附加到收缩列表的末尾。然后可以迭代此列表,直到所有中间收缩都完成。

步骤

首先,导入所需的库 −

import numpy as np

Tensors −

p = np.random.rand(2, 2)
q = np.random.rand(2, 5)
r = np.random.rand(5, 2)

要获取 einsum 表达式的最低成本收缩顺序,请使用 numpy.einsum+path() 方法 −

path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

显示路径信息 −

print(path_info[0])
print(path_info[1])

示例

import numpy as np
np.random.seed(123)

# 张量
p = np.random.rand(2, 2)
q = np.random.rand(2, 5)
r = np.random.rand(5, 2)

# 要获取 einsum 表达式的最低成本收缩顺序,请使用 Python 中的 numpy.einsum+path() 方法
path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

# 显示路径信息
print(path_info[0])
print(path_info[1])

输出

['einsum_path', (1, 2), (0, 1)]
  Complete contraction:  ij,jk,kl->il
         Naive scaling:  4
     Optimized scaling:  3
      Naive FLOP count:  1.200e+02
  Optimized FLOP count:  5.700e+01
   Theoretical speedup:  2.105
  Largest intermediate:  4.000e+00 elements
--------------------------------------------------------------------------
scaling                 current                   remaining
--------------------------------------------------------------------------
   3                 kl,jk->jl                  ij,jl->il
   3                 jl,ij->il                     il->il

相关文章


有用资源