Python 包提供收缩序优化、路径保存与加载、切片和数值收缩接口。
从本仓库构建 Python 包,请参阅源码安装。该方式不包含 Light/Heavy 引擎, 但可以执行已有收缩路径及给定的切片方案,并进行网络化简。 使用 Light/Heavy,需要安装包含引擎的完整 wheel,或单独配置兼容且已获授权的动态库。
从 ArcTN 1.0.0 下载页 获取完整 wheel,安装时不需要 Rust。根据平台选择说明 下载与系统、处理器和 Python 版本匹配的一份文件。完整包提供 CPython 3.12 和 3.13 版本; 源码支持 Python 3.9 及以上版本。 在该 wheel 文件所在的目录打开终端,然后运行:
python -m venv .venv
source .venv/bin/activate
python -m pip install --only-binary=arctn --find-links=. arctn==1.0.0arctn==1.0.0 指定安装版本;--find-links=. 让 pip 同时在当前目录查找兼容的 wheel,不需要手动填写文件名。
在 Windows PowerShell 中,请改用 .venv\Scripts\Activate.ps1 激活虚拟环境。
Linux、macOS 和 Windows 的 wheel 不能跨平台使用;具体可安装的组合以实际提供的文件为准。
完整 wheel 在 arctn/_lib 中包含编译后的 Light/Heavy 动态库。
导入 arctn 时会自动识别该动态库,调用 Light/Heavy 时加载,无需额外设置环境变量。
arctn_path、arctn_schedule、arctn_plan 等高层接口的参数和返回类型保持一致。
完整安装包供非商业使用;Light/Heavy 保持闭源,适用引擎二进制许可。 源码和动态库分别授权,说明见源码与 Light / Heavy。
NumPy 会作为依赖自动安装。如需与 Quimb 和 Cotengra 配合使用,请安装:
python -m pip install quimb cotengra opt_einsum从本仓库构建 Python 包需要 Python 3.9 及以上版本、Rust 1.83 及以上版本。 激活虚拟环境后,在仓库根目录运行:
python -m pip install "maturin>=1.9.3,<2"
python -m maturin develop --release --manifest-path pybind/Cargo.toml从本仓库源码构建的包不包含 Light/Heavy 动态库,但仍可执行给定的收缩路径与切片方案,
并进行网络化简。如需使用 Light 或 Heavy,需要提供兼容且已获授权的动态库,
并在首次调用搜索接口前,将 ARCTN_ENGINE_LIBRARY 设置为该库的绝对路径。
显式设置此变量时,会优先使用指定的动态库。
平台支持与 ABI 要求详见动态库接口说明。
| 接口 | 用途 |
|---|---|
arctn_path |
使用 Light 或 Heavy 搜索并返回收缩路径。 |
arctn_schedule |
返回最终路径、优化目标、指标,以及可选的切片索引。 |
arctn_plan |
返回 ArcTNExecutionPlan,保存网络结构、收缩路径与切片信息。 |
arctn_tree |
返回 Cotengra 的 ContractionTree。 |
arctn_contract |
优化收缩路径并执行数值收缩。 |
ArcTNOptimizer |
为 opt_einsum、Cotengra 和 Quimb 提供路径或收缩树搜索接口。 |
ArcTNCompiledContraction |
编译给定路径,以便重复执行收缩。 |
ArcTNExecutionPlan |
保存、加载和校验网络与收缩路径,并按其中的切片信息执行收缩。 |
arctn_simplify |
返回张量网络的结构化简结果。 |
使用 Light 或 Heavy 搜索时,接口会调用随包提供的动态库,或由 ARCTN_ENGINE_LIBRARY
显式指定的动态库。按给定路径执行收缩和化简张量网络不需要该动态库。
默认优化目标的权重为 flops_weight=1, read_write_weight=64。
设置 read_write_weight=0 可仅优化浮点运算量(FLOPs);
设置 flops_weight=0 可仅优化读写量。
target_size 是每个切片中生成的单个中间张量的元素数量上限,必须为正整数。
slicing_mode="fixed" 保持已选定的收缩路径;"dynamic" 还允许局部调整收缩路径。
请求切片时,两种模式都需要提供 target_size;
未提供 target_size 时,不能使用 "dynamic"。
arctn_schedule、arctn_contract、arctn_plan、arctn_tree 和 ArcTNOptimizer
均接受 allow_output_slicing,默认值为 False,只允许选择内部索引。
设为 True 后,Fixed 和 Dynamic 都可选择输出索引。执行时,内部切片在每个输出块内
求和,再按原输出轴顺序组装各块;保存并加载的计划也保留这一切片方案。
完整输出仍需全部存储空间,target_size 不限制最终输出大小或进程总内存。
外部数组后端在原后端上组装输出块,需要支持 stack;不会为组装而转成 NumPy 数组。
tnmpi 也能执行含输出索引切片的计划,每个进程保存完整输出,见 MPI 文档。
开启输出索引切片后,公开适配层先从 Light/Heavy 引擎取得未切片的收缩序,再调用本仓库的 公开切片代码。此流程兼容现有 ABI 1 引擎,不要求更换引擎动态库;详见 动态库接口说明。
将 ArcTNOptimizer(preset="heavy") 传给 Quimb 的 optimize 参数即可使用。
Quimb 负责构建张量网络并控制预处理;ArcTN 提供收缩路径,
也可通过 search() 返回包含切片索引的 Cotengra 收缩树。
实际的数值收缩仍由调用方选择的执行器完成。
如需直接执行数值收缩,arctn_contract 支持通过 backend="native"
使用 Rust CPU 执行器,也支持显式指定外部数组后端。
输入数组必须已属于所选后端;ArcTN 不会自动将数组传输到 GPU。