第三方工具接入

通过 Python 接口使用 ArcTN 进行收缩序优化,将路径或收缩树交给现有张量网络工具执行。

本页目录

作为 opt_einsum PathOptimizer

沿用快速开始中的三个数组,表达式 "ab,bc,cd->ad" 表示先把三个输入按 index 相连,最终保留 a、d。逗号分隔各个输入,-> 右侧指定结果的轴:

python
import opt_einsum as oe
from arctn import ArcTNOptimizer

equation = "ab,bc,cd->ad"
optimizer = ArcTNOptimizer(
    preset="heavy", seed=0, max_time=30,
    flops_weight=1, read_write_weight=64,
)
result = oe.contract(equation, *arrays, optimize=optimizer)

opt_einsum 直接调用 ArcTNOptimizer.__call__。PathOptimizer 接口只能返回收缩路径,不能同时返回切片 index。将 memory_limit 设为非 None 值,或在 ArcTNOptimizer 中设置了 target_size,都会使这个调用抛出 NotImplementedError。需要切片时,应通过 Cotengra/Quimb 调用 ArcTNOptimizer.search() 获取收缩树,或直接使用 arctn_tree / arctn_contract。

返回 Cotengra ContractionTree

沿用快速开始中的 inputs、output、size_dict 和 arrays。直接生成树并执行:

python
from arctn import arctn_tree

tree, info = arctn_tree(
    inputs, output, size_dict,
    preset="heavy", seed=0,
    target_size=2**24, slicing_mode="fixed",
    return_info=True,
)
result = tree.contract(arrays)

return_info=True 时返回 (tree, info),其中 info 是包含路径指标和切片信息的字典。若之前已用 arctn_plan 得到 plan,则不必再调用 arctn_tree 搜索,直接转换即可:

python
tree = plan.to_tree()
result = tree.contract(arrays)

ArcTNOptimizer.search() 和 arctn_tree() 会运行收缩序优化;plan.to_tree() 复用已有的收缩序和切片 index,不重新搜索。这些接口都可以得到 Cotengra 的 ContractionTree,用于表示收缩路径与切片信息。

在 Quimb 中使用

Quimb 通过 ArcTNOptimizer.search() 获取收缩树。将优化器传给 optimize,即可在电路任务或张量网络收缩中使用 ArcTN。

显式选择数组后端

使用 arctn_contract 时,可以选择 native 或显式指定数组后端。通过 Quimb 或 Cotengra 收缩树执行时,使用相应库的执行接口。