Skip to content

Latest commit

 

History

History

pybind

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

README.md

ArcTN Python 接口

Python 包提供收缩序优化、路径保存与加载、切片和数值收缩接口。

从本仓库构建 Python 包,请参阅源码安装。该方式不包含 Light/Heavy 引擎, 但可以执行已有收缩路径及给定的切片方案,并进行网络化简。 使用 Light/Heavy,需要安装包含引擎的完整 wheel,或单独配置兼容且已获授权的动态库。

安装完整 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.0

arctn==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 引擎,不要求更换引擎动态库;详见 动态库接口说明。

Quimb

将 ArcTNOptimizer(preset="heavy") 传给 Quimb 的 optimize 参数即可使用。 Quimb 负责构建张量网络并控制预处理;ArcTN 提供收缩路径, 也可通过 search() 返回包含切片索引的 Cotengra 收缩树。 实际的数值收缩仍由调用方选择的执行器完成。

如需直接执行数值收缩,arctn_contract 支持通过 backend="native" 使用 Rust CPU 执行器,也支持显式指定外部数组后端。 输入数组必须已属于所选后端;ArcTN 不会自动将数组传输到 GPU。

中文文档