量羽 QuantumQuill

Python API

按收缩序优化、保存与复用、数值执行和第三方接入查阅函数、参数与返回值。

本页目录

通用参数与分类

调用 Light / Heavy 的入口需要完整 Python wheel，或另行配置兼容的编译库，见安装。执行已有收缩序与网络化简不需要 Light / Heavy。

类别 API 返回内容
收缩序优化 arctn_path、arctn_schedule 路径，或包含路径与指标的报告
保存与加载 arctn_plan、ArcTNExecutionPlan 保存网络、收缩序和切片信息的对象
数值收缩 arctn_contract 计算结果，可选执行报告
编译与复用 ArcTNCompiledContraction 可重复调用的编译对象
第三方接入 arctn_tree、ArcTNOptimizer Cotengra 收缩树，或 ArcTNOptimizer 对象
网络化简 arctn_simplify 化简前缀与缩减网络信息

通用参数

以下参数只在对应函数签名列出时可用。签名中的 * 表示后续参数必须按名称传入。下面按签名列出参数和默认值；可直接运行的调用见快速开始和各示例。

参数 类型 含义
inputs Sequence[Sequence[Hashable]] 每个张量的 index 标签，内层顺序对应数组轴
output Sequence[Hashable] 保留的 index，顺序决定结果轴顺序
size_dict Mapping[Hashable, int] 已使用 index 的正整数维度，不得缺少或多出条目
preset "light" \| "heavy" 规划预设，默认 Heavy
seed int 基础随机种子；Light 和 Heavy 的额外独立起点种子由其派生，见单机并行
max_time float \| None 规划时间限制，单位秒；提供时必须有限且为正，在算法检查点生效
flops_weight float FLOPs 权重，有限且非负
read_write_weight float 读写量权重，有限且非负；两项权重不能同时为 0
target_size int \| None 每个切片中二元收缩结果的元素数上限；单张量网络检查最终输出
slicing_mode "fixed" \| "dynamic" Fixed 保持路径；Dynamic 允许调整路径，且必须提供 target_size
use_ssa bool 是否返回 SSA 编码；否则返回逐步重新编号的线性路径
rate_enabled bool 是否启用 rate 提前停止条件；关闭它不会关闭其他停止条件
backend str "native" 或显式数组后端名，不接受 "auto"
return_info bool 是否同时返回信息字典

网络输入规则见张量网络输入。target_size 的单位是元素，不是字节、RSS 或显存。数组后端和异步计时见执行后端。

收缩序优化

arctn_path

python
arctn_path(
    inputs,
    output,
    size_dict,
    *,
    preset="heavy",
    seed=0,
    use_ssa=False,
    max_time=None,
    flops_weight=1.0,
    read_write_weight=64.0,
)

返回 list[tuple[int, int]]。只寻路，不执行数组。默认返回线性路径；use_ssa=True 返回 SSA 路径。此函数不接收切片参数，因为路径列表不能同时携带切片 index。

arctn_schedule

python
arctn_schedule(
    inputs,
    output,
    size_dict,
    *,
    preset="heavy",
    seed=0,
    target_size=None,
    slicing_mode="fixed",
    max_time=None,
    use_ssa=False,
    flops_weight=1.0,
    read_write_weight=64.0,
    rate_enabled=True,
)

返回 dict，包含路径、目标权重、路径指标、编译库记录的搜索耗时 wall_s 和可选切片信息。不读取或执行数值数组，不返回内部候选和阶段记录。字段含义及示例见使用 Heavy：读取返回结果。

保存与加载

arctn_plan

python
arctn_plan(
    inputs,
    output,
    size_dict,
    *,
    preset="heavy",
    seed=0,
    target_size=None,
    slicing_mode="fixed",
    max_time=None,
    flops_weight=1.0,
    read_write_weight=64.0,
    rate_enabled=True,
)

返回 ArcTNExecutionPlan。只规划，不执行数组。该对象保存网络、SSA 路径和切片信息，不保存数值数组、设备或已编译的计算内核。

ArcTNExecutionPlan

ArcTNExecutionPlan 可保存为文件，在其他进程中加载；也可转换为收缩树，或在无切片时编译后重复执行。完整示例见保存与复用收缩序。

从报告创建

python
ArcTNExecutionPlan.from_schedule(
    report,
    *,
    inputs,
    output,
    size_dict,
)

report 必须包含 SSA 路径，即用 arctn_schedule(..., use_ssa=True) 生成。三个网络参数均为必填项。返回 ArcTNExecutionPlan，不再次寻路。

从文件加载

python
ArcTNExecutionPlan.load(path)

读取并验证包含完整网络的 version 2 JSON 文件，返回 ArcTNExecutionPlan。验证失败抛出 PlanValidationError；文件不存在或无法读取时，仍可能抛出文件系统异常。

从字典加载

python
ArcTNExecutionPlan.from_dict(artifact)

从已解析的字典加载，使用与文件加载相同的验证规则。

保存文件

python
plan.save(path)

以 JSON 保存网络、收缩序和切片信息，采用原子写入；不保存数组或进程内编译状态。

导出字典

python
plan.to_dict()

返回包含网络、收缩序和切片信息的可序列化字典。

检查网络或数组

python
plan.validate(
    inputs=None,
    output=None,
    size_dict=None,
    *,
    arrays=None,
)

三个网络参数要么全部省略，要么一起提供。可以另传 arrays 检查数组数量、形状和数据类型。成功返回 None，失败抛出 PlanValidationError。

转换为线性路径

python
plan.to_linear_path()

返回 opt_einsum 使用的线性路径，不重新规划。该列表本身不携带切片信息。

转换为收缩树

python
plan.to_tree()

返回 Cotengra ContractionTree，包含已有路径和 sliced indices。需要安装 Cotengra，不重新寻路。

编译收缩序

python
plan.compile(
    *,
    backend="native",
)

返回 ArcTNCompiledContraction。仅在实际切片集合为空时可用；有切片时应使用 plan.execute()。

按保存的收缩序执行

python
plan.execute(
    arrays,
    *,
    backend="native",
    return_info=False,
)

沿保存的路径和切片方案执行，不重新规划。默认返回结果数组；return_info=True 返回 (result, info)。信息字段随执行方式而异，不应假设与 arctn_contract 的计时字段完全相同。

数值收缩

arctn_contract

python
arctn_contract(
    inputs,
    output,
    size_dict,
    arrays,
    *,
    preset="heavy",
    seed=0,
    target_size=None,
    slicing_mode="fixed",
    max_time=None,
    backend="native",
    flops_weight=1.0,
    read_write_weight=64.0,
    return_info=False,
)

返回 结果数组，或 (result, info)。原生后端返回 NumPy 数组；外部后端返回相应数组库的对象。所有输入须使用同一种受支持的数据类型：float32、float64、complex64 或 complex128。

外部后端可能异步执行。主机函数返回不代表设备计算完成；GPU 耗时需要显式同步后测量。execution_wall_s 只对 native 和 numpy 提供，其他后端为 None。完整字段见执行后端。

编译与复用

ArcTNCompiledContraction

按固定网络和 SSA 收缩序预先准备执行：原生后端准备轴变换与矩阵乘法参数，外部后端编译 opt_einsum 表达式。该对象只在当前进程中使用，不能保存为文件。

创建编译对象

python
ArcTNCompiledContraction.compile(
    inputs,
    output,
    size_dict,
    *,
    ssa_path,
    backend="native",
)

ssa_path 必须完整合法。编译不寻路，也不支持切片集合。

执行与统计

python
result = compiled.execute(arrays)
info = compiled.stats()

直接调用 compiled(arrays) 与 execute 等价。多次调用使用同一网络和路径；每次调用仍验证输入。示例见编译并重复执行。

第三方接入

arctn_tree

python
arctn_tree(
    inputs,
    output,
    size_dict,
    *,
    preset="heavy",
    seed=0,
    target_size=None,
    slicing_mode="fixed",
    max_time=None,
    flops_weight=1.0,
    read_write_weight=64.0,
    return_info=False,
)

返回 Cotengra ContractionTree，或 (tree, info)。需要 Cotengra 和 opt_einsum。树携带收缩路径及切片 index，可供 Quimb / Cotengra 执行。

ArcTNOptimizer

python
ArcTNOptimizer(
    *,
    preset="heavy",
    seed=0,
    max_time=None,
    flops_weight=1.0,
    read_write_weight=64.0,
    target_size=None,
    slicing_mode="fixed",
)

ArcTNOptimizer 将 ArcTN 的收缩序优化接入两类接口：opt_einsum 的路径接口与 Cotengra / Quimb 的收缩树接口。使用示例见在 Quimb 中使用。

返回路径

python
optimizer(inputs, output, size_dict, memory_limit=None)

对应 ArcTNOptimizer.__call__，返回线性路径。由于列表无法表达切片集合，提供 memory_limit 或在 ArcTNOptimizer 中配置 target_size 后使用此调用，均会抛出 NotImplementedError。

python
optimizer.search(inputs, output, size_dict)

对应 ArcTNOptimizer.search，委托 arctn_tree 返回含切片信息的 ContractionTree。

收缩序优化与网络化简

ArcTNOptimizer 用于收缩序优化；只需要化简结构信息时，使用下面的 arctn_simplify。

网络化简

arctn_simplify

python
arctn_simplify(inputs, output, size_dict)

返回 化简统计、前缀和缩减网络信息。不执行 Light / Heavy 搜索，也不执行数组计算。规则见网络化简。

错误处理

常见异常包括 TypeError、ValueError、PlanValidationError、NotImplementedError 和缺少可选依赖时的 ImportError。具体例子见检查输入与处理错误。

下载原始 Markdown 文档 ↓