数值收缩
沿收缩序计算数组,使用原生 Rust CPU 后端或其他数组库,并按 output 指定的轴顺序返回结果。
本页目录
从收缩序到数值结果
收缩序确定计算的先后顺序;数值执行则读取输入数组,逐步计算中间张量和最终结果。
| 用法 | 入口 |
|---|---|
| 搜索收缩序并执行,可同时切片 | arctn_contract |
| 执行已保存的收缩序和切片集合 | ArcTNExecutionPlan.execute |
| 在 Rust 中执行给定 SSA path | contract_network |
| 在 Rust 中执行给定 SSA path 和切片集合 | contract_network_sliced |
Python 默认使用 backend="native",即 ArcTN 的 Rust CPU 执行器。也可指定 NumPy、CuPy 等执行后端。使用 Quimb 时,将 ArcTNOptimizer 传给 optimize=,由 Quimb/Cotengra 完成数值收缩,见 Quimb 示例。
原生二元收缩步骤
原生执行先检查数组数量、各轴的维度和实际存储的元素数是否与网络一致,再检查收缩路径能否从这些输入得到最终结果。例如,形状为 (2, 3) 的输入应有 6 个元素,不能按 (3, 2) 的轴顺序直接代入。具体例子见输入和路径检查。
检查通过后,每一步按“哪些 index 需要保留、哪些需要求和”排列输入轴,再把同类轴合并为矩阵维度,执行矩阵乘法(GEMM)。一个共享 index 若还出现在后续张量或最终输出中,此时必须保留;只有不再需要它时才沿该 index 求和。
C[b,i,j] = Σk A[b,i,k] B[b,k,j]
- b
- 两侧共享,保留为 batch 轴
- k
- 两侧共享,本步求和消去
- i、j
- 分别来自 A、B 的自由轴
预处理后,对轴进行排列和合并,即可转为 batched GEMM;结果再恢复为输出张量的轴顺序。
两个输入共享且收缩后仍需保留的 index 对应批次轴;共享但需要求和的 index 对应矩阵乘法的收缩轴。只在一侧保留的 index 分别组成矩阵的行、列轴。执行完成后,结果按 output 指定的顺序排列轴。
输入、dtype 与结果
| Rust 类型 | Python 数据类型 |
|---|---|
f32 |
float32 |
f64 |
float64 |
Complex32 |
complex64 |
Complex64 |
complex128 |
每次执行的数组数量和形状必须匹配网络,全部输入使用同一种受支持的数据类型。Python 原生后端将输入整理为 C 连续存储的 NumPy 数组;结果轴顺序由 output 指定。
执行采用稠密二元收缩,不做键维度截断或低秩近似。PathStats 中名为 FLOPs 的指标,统计的是根据网络结构和收缩路径计算的标量乘法次数;它不能直接等同于 CPU 指令数、内存流量或实际运行时间。
两条路径即使 PathStats 相同,也可能因轴排列、转置和 GEMM 形状不同而有不同运行时间。后端性能需要在相同路径、输入、dtype、线程和计时边界下单独测量。
直接执行
Rust:沿已有路径执行,不重新寻路
use arctn::contract_network;
let result = contract_network(&net, tensors, &ssa_path)?;
Python:寻路并执行
from arctn import arctn_contract
result, info = arctn_contract(
inputs, output, size_dict, arrays,
preset="heavy", backend="native",
return_info=True,
)
info 记录实际使用的后端和执行时间。若只改变数组值、网络结构和收缩序保持不变,可以编译一次并重复执行。需要跨进程保存收缩序和切片集合时,使用 arctn_plan,见保存与复用收缩序。