张量网络与收缩路径

用 index 和维度描述张量网络，通过收缩树与路径表示收缩顺序，并检查输入和路径是否合法。

本页目录

张量网络的表示

张量网络指定哪些张量相乘、沿哪些 index 求和。这里的 index（索引）用于标记张量的轴；矩阵的行和列就是两个轴，各轴的长度称为维度。例如，三个矩阵的乘积可以写成

Y_{ad}=\sum_b\sum_c A_{ab}B_{bc}C_{cd}.

b、c 是求和 index，a、d 保留为结果的两个轴。ArcTN 根据这些连接关系和维度优化收缩序，不需要读取张量数值。

点表示张量，连线表示共享 index；两端的 a、d 是输出 index。

写成 ArcTN 输入

设 $A$ 、 $B$ 、 $C$ 的形状分别为 $2\times3$ 、 $3\times4$ 、 $4\times5$ ：

python

inputs = [("a", "b"), ("b", "c"), ("c", "d")]
output = ("a", "d")
size_dict = {"a": 2, "b": 3, "c": 4, "d": 5}

inputs 的外层顺序对应输入数组的顺序：本例依次为 A、B、C，执行时也按 [A, B, C] 传入数组。每个元组再按该数组的轴顺序列出 index；例如 ("a", "b") 表示第一个矩阵的行对应 a、列对应 b。它不指定收缩序。
output 指定结果轴；本例结果矩阵的行对应 a、列对应 d，形状为 (2, 5)。
size_dict 给出每个 index 的维度。相同标签表示同一个 index，与标签本身的名称无关。

Python 接口接受可哈希标签，在进入 Rust 前转换为整数编号。若使用 Quimb，可直接传入 ArcTNOptimizer，由 Quimb 提供这些网络信息，见Quimb 示例。

普通边、超边与输出 index

一个 index 连接两个不同张量时，可以画成普通边；连接三个或更多不同张量时，是超边（hyperedge）。未列入 output 的 index 最终求和消去，输出 index 则保留。

同一张量的两个轴也可以共用一个标签。如果这个 index 不在其他张量和输出中出现，则取迹，例如 $\sum_i A_{ii}$ ；否则先取对角，保留一个 index 参与后续收缩。

同一个 index 在单个输入张量中最多出现两次；出现三次或更多时，网络校验会返回错误。

收缩树与路径

网络给出了要计算的张量表达式，但没有指定先收缩哪些张量。对于上面的例子，可以先计算 $AB$ ，也可以先计算 $BC$ ：

Y=(AB)C=A(BC).

两种顺序在数学上得到相同结果，但计算量和中间张量大小不同：

收缩序	第一步结果	总标量乘法次数	最大收缩结果（元素）
$(AB)C$	$2\times4$	$2\cdot3\cdot4+2\cdot4\cdot5=64$	10
$A(BC)$	$3\times5$	$3\cdot4\cdot5+2\cdot3\cdot5=90$	15

最大收缩结果包括最终输出；第一条路径虽然先产生 8 个元素，最终输出仍有 10 个元素。

括号结构对应一棵二叉收缩树：叶节点是输入，内部节点是二元收缩，根是最终结果。路径则按执行顺序列出这些步骤。同一棵树中互不依赖的步骤可以交换先后顺序。

SSA path：固定张量编号

ArcTN 内部使用 static single assignment（静态单赋值，SSA）path。这里的含义是：每个输入和每步结果各有一个固定编号，每个新结果分配一个新编号。若有 n 个输入，输入编号为 0..n-1；第 s 步的结果编号固定为 n+s，其中 s 从 0 开始。参与收缩的两个输入不再用于后续步骤，新结果使用新编号。

路径中的整数指向整个张量。它与 inputs 中标记张量轴的 "a"、"b" 等 index 标签是两回事。

以上述三张量网络的 $A(BC)$ 顺序为例：

text

输入编号：A = 0，B = 1，C = 2
SSA path：[(1, 2), (0, 3)]

第 0 步：B 与 C 收缩，结果编号为 3
第 1 步：A 与结果 3 收缩，最终结果编号为 4

Linear path：当前列表位置

opt_einsum 的 linear path 使用当前张量列表中的位置。每收缩一步，从列表中移除本步的两个输入张量，再把新结果追加到末尾，因此后续位置会变化。上例第一步之后，列表由 [A, B, C] 变成 [A, BC]，第二步使用位置 (0, 1)；完整 linear path 为 [(1, 2), (0, 1)]。

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)

SSA path 和 linear path 都是整数对列表，但编号含义不同。保存或传递路径时，需要同时说明路径格式。

候选收缩树的比较

同一个张量网络可以有不同的收缩树。搜索方法产生并改进候选，再按完整路径的目标值比较。下图用四个输入张量 A、B、C、D 展示五种括号结构：

P₁

P₂

P₃

P₄本例的最优候选

P₅

五棵树表示同一网络的不同候选，本例假设 P₄ 的目标值最低。数量和树形仅作示意，不表示 Light 或 Heavy 固定搜索五个候选。

按目标函数选择

设 $\mathcal C$ 为本次搜索中参与最终比较的完整候选路径集合， $J(P)$ 为路径 $P$ 的目标值，则选择

P^* \in \operatorname*{arg\,min}_{P\in\mathcal C} J(P).

$P^*$ 是目标值最低的候选之一。FLOPs 与读写量如何组成 $J(P)$ ，见路径指标与优化目标。

图中比较未切片路径。切片后按全部切片的总工作量比较；动态切片还可能调整路径，见切片原理。

最佳候选的含义

Light 和 Heavy 返回本次搜索中找到的最佳候选，不保证局部最优或全局最优。simulate_path 可独立检查路径是否合法。

并行搜索与随机种子

RAYON_NUM_THREADS 设置搜索使用的 Rayon 线程池宽度，seed 设置本次搜索的基础随机种子。普通调用只需执行一次 Light 或 Heavy 搜索；线程设置与多次重复测量的区别见使用 Heavy与单机并行。

输入校验

输入校验先确认网络是否定义完整：至少有一个输入张量，每个 index 都有正整数维度，output 只包含输入中出现的 index，且不重复。size_dict 也不能多出未使用的条目。例如，上面的网络若漏写 "b": 3，就无法确定 A 与 B 相连的维度。

执行时再检查数组数量和形状。例如，inputs[0] = ("a", "b")、size_dict["a"] = 2、size_dict["b"] = 3 要求第一个数组的形状为 (2, 3)。维度与实际数组不符、或张量元素数超过平台可表示的范围时，接口会返回错误。

直接使用 Rust 结构体字段构造 TensorNetwork 时，应显式调用 validate()。

路径校验

simulate_path 按 SSA 路径检查每一步是否合法，并计算 PathStats。它只处理网络结构和维度，不执行实际数组收缩。指标定义和调用示例见路径指标与优化目标。

每一步必须引用两个不同且仍可使用的张量编号；最终只剩一个结果，其 index 与 output 一致。单张量网络可以使用空路径，必要的求和、取迹或取对角由一元操作处理。错误处理见检查输入与处理错误。

下载原始 Markdown 文档 ↓