收缩路径

收缩路径规定二元收缩的括号结构;ArcTN 内部使用 SSA path,并与 opt_einsum linear path 互相转换。

本页目录

路径的作用

同一个张量网络可以按不同顺序完成相同的数学收缩,但中间张量大小、FLOPs 和逻辑读写量可能相差很大。收缩路径把多张量表达式分解成一系列二元步骤,也等价地定义了一棵二叉收缩树。

同一棵二叉收缩树可以写成稳定的 SSA path,也可以转换为每步重新编号的 recycled linear path。

局部选择与完整路径

路径优化比较的是完整合法路径。生成器的局部启发式代价与整条路径的最终规划目标是两个不同层次。

SSA path 的编号规则

ArcTN 的内部表示是 static single assignment(SSA)path。若有 n 个输入,输入编号为 0..n-1;第 s 步的结果编号固定为 n+s。两个 operand 在该步被消费,之后不能再次引用。

text
四个输入:0, 1, 2, 3
SSA path:[(0, 1), (2, 3), (4, 5)]

step 0: (0, 1) -> 4
step 1: (2, 3) -> 5
step 2: (4, 5) -> 6  # 最终根结果
规则 校验结果
一步引用同一个编号两次 自收缩,路径无效
引用不存在或已消费的编号 路径无效
有 n 个输入但没有留下一个完整根结果 路径不完整
最终腿与 output 不一致 路径无效

SSA path 与 recycled linear path

opt_einsum 常用 recycled linear path:整数表示当前存活 operand 列表中的位置。每收缩一步,列表删除两个 operand,再把新结果追加到列表,因此后续位置会变化。SSA 编号则在整条路径中保持不变。

表示 整数含义 主要用途
SSA path 固定的输入或中间结果编号 ArcTN 内部搜索、稳定存储、重放和 Rust 接口
recycled linear path 当前存活 operand 列表的位置 opt_einsum PathOptimizer 与 Python 生态
可变收缩树 父子节点与局部树结构 旋转、子树重构、退火和回火
python
from arctn import arctn_path

linear_path = arctn_path(inputs, output, size_dict)
ssa_path = arctn_path(
    inputs, output, size_dict, use_ssa=True
)

Warning

最常见的路径格式错误

两种路径看起来都是整数对列表,但坐标系不同。保存、加载或传给编译器时必须同时记录路径格式,不能靠列表外观猜测。

重放路径而不是执行张量

simulate_path 按 SSA 顺序重建每一步的腿集合,检查路径是否合法,并计算统一的 PathStats。它只处理结构和维度,不分配或收缩实际张量数组。

rust
use arctn::simulate_path;

let stats = simulate_path(&net, &ssa_path)?;
println!("log10 FLOPs = {}", stats.log10_flops);
println!("log2 max result = {}", stats.log2_max_size);

重放时,每个二元步骤引用的两个张量必须仍然存活;非空路径完成后应只剩一个结果,且其腿集合与 output 一致。单张量网络可使用空路径,其剩余求和由一元操作处理。外部生成的路径也可以用同一接口检查和计数。