数值收缩

沿收缩序计算数组,使用原生 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]

A[b,i,k]batch × m × k
B[b,k,j]batch × k × n
C[b,i,j]batch × m × n
b
两侧共享,保留为 batch 轴
k
两侧共享,本步求和消去
i、j
分别来自 A、B 的自由轴

预处理后,对轴进行排列和合并,即可转为 batched GEMM;结果再恢复为输出张量的轴顺序。

二元收缩转为 batched matrix multiplication 的例子:b 保留,k 求和,i、j 组成结果的自由轴。多个同类轴可排列并合并为一个矩阵维度。

两个输入共享且收缩后仍需保留的 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:沿已有路径执行,不重新寻路

rust
use arctn::contract_network;

let result = contract_network(&net, tensors, &ssa_path)?;

Python:寻路并执行

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,见保存与复用收缩序。