使用 Quimb 构建并执行张量网络

让 Quimb 负责量子电路任务或自定义 TensorNetwork,让 ArcTN 负责路径规划与可选切片。

本页目录

组件职责

本教程展开的是“接入现有生态”路线:Quimb 负责把任务变成待收缩张量网络,ArcTN 负责路径规划和可选切片,Quimb/Cotengra 负责按返回的收缩树执行。若已经直接持有结构和数组,也可以改用 arctn_contract(..., backend="native"),不经过 Quimb/Cotengra。

组件 负责 不负责
Quimb 电路和张量网络构造、任务化简、数值 contraction、数组后端 替 ArcTN 决定其内部搜索策略
ArcTN 收缩路径、完整路径指标,以及通过 ArcTNOptimizer.search 或 arctn_tree 请求的切片腿 量子门语义、振幅定义或采样逻辑
Cotengra 承载 ContractionTree,并执行 tree-aware contraction ArcTN 的 Light / Heavy 搜索

这条接口不仅适用于量子电路。只要 Quimb TensorNetwork 满足 ArcTN 当前的输入约束,就可以把 ArcTNOptimizer 传给 optimize=:同名 index 的维度必须一致,output 不能重复,而且同一个 index 在单个输入张量内至多出现两次。

安装

bash
git clone https://github.com/2hengy1/arctn.git
cd arctn/pybind
python -m pip install -e '.[quimb]'

从源码安装完整 Quimb 集成时,quimb extra 会安装 Quimb、Cotengra 和 opt_einsum。若使用 CuPy、PyTorch 或 JAX,还需要另行安装与本机硬件和驱动匹配的数组库。

从 Quimb Circuit 计算一个振幅

python
import numpy as np
import quimb.tensor as qtn
from arctn import ArcTNOptimizer

n = 8
rng = np.random.default_rng(0)
circ = qtn.Circuit(n)
for depth in range(6):
    for q in range(n):
        circ.apply_gate("RY", float(rng.uniform(0, 2 * np.pi)), q)
    for q in range(depth % 2, n - 1, 2):
        circ.apply_gate("CNOT", q, q + 1)

optimizer = ArcTNOptimizer(preset="light", seed=0)
amplitude = circ.amplitude(
    "0" * n,
    optimize=optimizer,
    simplify_sequence="ADCRS",
)
print(amplitude)

调用 amplitude 时,Quimb 先把振幅任务变成张量网络并执行局部化简,然后把网络的 inputsoutputsize_dict 交给 ArcTN。对于非平凡网络,Cotengra/Quimb 会调用 ArcTNOptimizer.search();ArcTN 返回 ContractionTree,随后由 Quimb/Cotengra 执行。平凡收缩可能绕过优化器,详见下文。

Quimb Circuit 任务 ArcTN 的作用 结果语义
amplitude 规划振幅网络 给定 bitstring 的复振幅
local_expectation 规划局域期望值网络 局域算符期望值
compute_marginal / partial_trace 规划边缘概率或约化态网络 数组或密度矩阵
sample 为条件采样过程中产生的网络规划路径 指定数量的 bitstrings
to_dense 规划完整态矢量网络 稠密态矢量

路径按网络几何复用

不同任务生成的网络不同,同一个 Circuit 不等于只需要一条可供所有任务复用的路径。是否能够复用,取决于待收缩网络的 inputs、output 和尺寸是否相同。

理解 Quimb 的任务化简

Quimb 的 amplitudelocal_expectationpartial_tracecompute_marginal 和常用 sampling 入口当前默认使用 simplify_sequence="ADCRS"to_dense 当前默认使用 R。这些都是 Quimb 在 ArcTN 寻路之前执行的局部张量网络化简,因此 ArcTN 看到的是化简后的 residual tensor network,而不是原始门列表。下表解释本教程振幅示例使用的 ADCRS

字母 Quimb 方法 作用
A antidiag_gauge 翻转与反对角结构相连的 bond gauge,为后续 diagonal reduction 创造条件
D diagonal_reduce 折叠对角张量的冗余轴;可能产生由多个张量共享的 hyperedge
C column_reduce 识别某个指标只有一个取值使张量非零的情况,并将该指标固定为这个取值
R rank_simplify 收缩不会增大张量阶数(指标数)的相邻张量;该选择规则只需要网络结构
S split_simplify 在给定数值容差内寻找低秩分解并拆分张量

Quimb 按给定字母顺序循环这些步骤,直到张量数和 index 数不再减少。它与 ArcTN 内部可能执行的结构处理不是同一个 pass;ArcTN 从 Quimb 交来的 residual tensor network 开始规划。

python
rehearsal = circ.amplitude_rehearse(
    b="0" * n,
    simplify_sequence="ADCRS",
    optimize=optimizer,
)
tn = rehearsal["tn"]
tree = rehearsal["tree"]
print(len(tn), tree.total_flops())

Normalization exponent

在 Quimb 内完成建网和收缩时,normalization exponent 已由 Quimb 管理;再次手工乘缩放因子会造成重复缩放。

通过 ArcTNOptimizer 直接切片

python
from arctn import ArcTNOptimizer

tn = circ.amplitude_tn(
    "0" * n, simplify_sequence="ADCRS"
)
output = tuple(tn.outer_inds())

optimizer = ArcTNOptimizer(
    preset="heavy", seed=0,
    target_size=2**28,
    slicing_mode="fixed",
)
value = tn.contract(
    all, output_inds=output, optimize=optimizer,
)

对于这条调用链中的非平凡收缩,Quimb/Cotengra 会调用 ArcTNOptimizer.search()。返回的 ContractionTree 同时携带 contraction path 和 sliced indices;Quimb/Cotengra 随后执行各个 slice 并求和,Quimb 则应用网络的 normalization exponent 所记录的缩放因子。fixed 保持当前 objective 下选出的最佳未切片候选路径不变;dynamic 允许切片后做局部路径重构。平凡网络可能绕过优化器,详见下文。

target_size 的单位

target_size 是每个 slice 中二元收缩结果元素数的硬约束;单张量网络则检查最终输出。它不限制全部临时张量,也不是字节数、RSS、VRAM 或整个进程的内存上限。

同一个 ArcTNOptimizer 同时支持两种上层协议:opt_einsum 直接调用 __call__ 时只能返回 linear path;Cotengra/Quimb 调用 search() 时返回 ContractionTree,因此可以保留 sliced indices。若调用方需要先取得树再决定何时执行,也可以显式使用 arctn_tree

平凡网络的调用边界

Cotengra 对只有一到两个输入张量的平凡收缩可能直接构造树,不调用外部 optimizer。必须对这类网络严格施加 target_size 时,应直接使用 arctn_treearctn_contract

手工组装任意 TensorNetwork

python
import numpy as np
import quimb.tensor as qtn
from arctn import ArcTNOptimizer

rng = np.random.default_rng(0)
ta = qtn.Tensor(rng.normal(size=(2, 3)), inds=("a", "x"), tags={"A"})
tb = qtn.Tensor(rng.normal(size=(3, 4)), inds=("x", "y"), tags={"B"})
tc = qtn.Tensor(rng.normal(size=(4, 5)), inds=("y", "z"), tags={"C"})
td = qtn.Tensor(rng.normal(size=(5, 2)), inds=("z", "b"), tags={"D"})
tn = qtn.TensorNetwork([ta, tb, tc, td])

result = tn.contract(
    all,
    output_inds=("a", "b"),
    optimize=ArcTNOptimizer(preset="light", seed=0),
)
array = result.data
print(array.shape)  # (2, 2)

相同的 index 名表示连接:xyz 在收缩时求和;output_inds=("a", "b") 保留两个开放指标,并规定结果轴的顺序。每个同名 index 在所有张量上的维度必须一致。

高级用法:拆出网络并交给 arctn_contract

python
from arctn import arctn_contract

tensors = list(tn)
inputs = [tuple(tensor.inds) for tensor in tensors]
output = ("a", "b")
size_dict = dict(tn.ind_sizes())
arrays = [tensor.data for tensor in tensors]

value = arctn_contract(
    inputs, output, size_dict, arrays,
    preset="heavy", backend="numpy",
)
exponent = float(getattr(tn, "exponent", 0.0) or 0.0)
value = value * (10.0 ** exponent)

这条路线适合需要 ArcTN native executor、显式外部后端,或完整 return_info 报告的场景。若网络曾由 Quimb 在 equalize-norms 模式下化简,tn.exponent 保存的是独立缩放因子的十进制指数。提取数组并在 Quimb 之外完成收缩后,将结果乘以一次 10.0 ** tn.exponent

实验和性能比较边界

比较对象 固定条件 独立记录项
路径质量 同一个化简后 TensorNetwork、同一 objective 寻路时间、FLOPs、read/write、最大中间张量
端到端任务 同一 Circuit 任务、化简序列、dtype 和 backend 建网、化简、寻路、执行和同步后的总时间
切片执行 同一 target_size 与结果精度 slice 数、每片峰值、总 FLOPs、设备完成时间

GPU 后端可能异步执行;只测 Python 调用返回时间会漏掉设备工作。比较 CuPy 或 JAX 时,应在计时终点执行对应后端的同步操作。