RSTSR v0.10：Rust 原生、全面兼容 Python Array API 标准的张量库

RSTSR 历时两年，现在到了阶段性的里程碑：CPU 下高效的、原生 Rust 的、与 NumPy 接口高度兼容的 n-维张量科学计算库。

我们认为，NumPy 在高维张量的接口设计是非常优秀的。我们在 RSTSR 中证明，Rust 语言原生科学计算可以做到 写起来、读起来像 NumPy，同时达到更高效、更安全、更适合并行计算。

• 项目地址：https://github.com/RESTGroup/rstsr
• API 文档：https://docs.rs/rstsr/latest/rstsr

详见 RSTSR v0.10 相关博客：https://restgroup.github.io/rstsr-book/zh-hans/blog/rstsr-third-report

对于 NumPy 用户，RSTSR 提供了 NumPy-RSTSR Cheatsheet 让你快速熟悉 RSTSR。

全面兼容 Python Array API 2025 标准

RSTSR 作为原生的 Rust 库，它实现了几乎所有 Python Array API 2025 标准的接口。RSTSR 不再是“受 Array API 启发”，而是真正实现了 Array API 标准的非 Python 程序库。

Python Array API 2025 标准涵盖了几乎所有 NumPy 常用与重要的功能，包括：

• 张量的创建：zeros, ones, full, arange, linspace, eye, meshgrid 等；
• 张量的索引：基础索引 (basic indexing/slicing)、高级索引 (advanced/fancy indexing)、元素索引 (element indexing)；
• 张量的形状操作：reshape, transpose, swapaxes, moveaxis, expand_dims, squeeze 等；
• 张量的广播操作：broadcast_to, broadcast_arrays 等；广播机制也应用于张量的算数运算等其他操作；
• 张量的算术运算：加减乘除、三角函数等；
• 逻辑运算、比较运算等二元操作；
• 张量的线性代数运算：matmul, vecdot, eigh, svd 等；
• 张量的统计运算：sum, mean, var, std, min, max 等；
• 张量的其他操作：concatenate, stack, split, repeat, tile 等。

基于 RSTSR v0.10.0、Faer 后端的的简单 Python 封装程序 rstsr-faer-py 可以在 PyPI 上下载；对 Array API 2025 测试套件 结果是 1336 Pass / 17 Fail / 33 Skip (对比 NumPy 2.5 的结果是 1339 Pass / 42 Fail / 5 Skip)。其中，14 Fail / 33 Skip 是测试套件自身、以及 RSTSR Faer 后端未实现 FFT 功能；其余 3 Fail 是实际应用情景中罕见的 SVD 浮点数过大溢出问题。可以表明 RSTSR 除 FFT 功能外，正确实现了所有 Array API 2025 的接口。

尽管我们并非真的实现了所有 Array API 2025 的接口，但那是因为少量接口由于 Rust 语言本身的特性难以或不可能实现。因此我们即使在 Python 下存在解决方案，但原生 Rust 一一对应地实现。这些被排除的功能包括

• Python 中允许的 __ipow__, __imatmul__，但 Rust 不存在 (但可以用函数实现的) 的运算符重载；
• Python 允许的动态类型列表，但 Rust 通过静态类型实现非常困难；受到影响的函数包括 concat, stack (但只要是相同类型的张量，RSTSR 仍然可以使用 rt::concat 和 rt::stack)。

RSTSR 的主要特性

• 使用 % 作为矩阵乘法算符：这是 RSTSR 的一个特色。你可以使用 &a % &b 来执行 a.matmul(&b)，这使得矩阵乘法的语法更加简洁。
• 优秀的并行管理：针对多线程 BLAS 后端，我们通过 rayon 的线程控制，可以保证无论是否在 rayon 并行线程内部还是外部调用 RSTSR 的通用计算或矩阵线性代数，都能充分利用并行设备优势、且不会超出用户所指定的线程数。
• 同时支持 row-major 与 column-major：RSTSR 默认支持 row-major 内存布局；但我们同时允许 column-major 内存布局，且在现实项目 REST 电子结构程序中已经大量使用。
• 支持多种 BLAS 后端：RSTSR 可以使用 Faer 默认多线程后端，也可以使用 OpenBLAS、Intel MKL、BLIS、KML 等多种 BLAS 后端；在 CPU 下的性能表现与 NumPy 相当。支持如 solve, svd, eigh 等矩阵线性代数运算。
• 可扩展后端：我们采用泛型以支持多后端。用户不需要修改 RSTSR 的源代码，不需要要求数据在内存 (GPU/硬盘/网络都可以承载数据)，只要实现新的后端 trait 接口，就可以使用自定义后端实现数学计算。
• 支持多种数据类型：RSTSR 支持通常浮点数，也支持任意数据结构 (包括字符串等需要 alloc 的类型)；在电子结构程序中，f64 与 c64 是最常用的数据类型，我们对这些数据类型进行了支持和特殊优化。
• 零开销 reshape 支持。在合理的张量布局下 (不一定必须是内存完全连续的)，RSTSR 可以使用视窗 (view) 的方式实现 reshape，而无需额外的内存开销。
• Einstein Summation 支持。在外接 TBLIS 张量缩并库的情况下，RSTSR 可以使用与 NumPy 相同的 Einstein Summation 语法进行张量缩并计算。详细见 rt::tblis::einsum。
• 高效的 Vecdot 支持。对于电子结构，特别是一部分稠密密度泛函矩阵计算，"pi, pi -> p" 类型的张量缩并非常常见，且确实是不小的内存带宽耗时瓶颈。我们对 rt::vecdot 函数的实现并非是 naive 的；它在多核并行下性能优秀，比单纯的矩阵数乘后再求和要快许多，且内存开销更小。

RSTSR 在电子结构实际项目中的应用

高维张量库 RSTSR 的设计初衷，是更好地解决电子结构计算中高维张量的存储与计算。RSTSR 基于 REST (Rust-based Electronic Structure Toolkit) 电子结构程序平台开发。

RSTSR 的实战典型示例是，我们在 REST 电子结构程序中，实现了高效的杂化泛函二阶解析梯度 (REST PR !154)。

Hessian 二阶解析梯度可以给出重要的化学性质，包括分子振动频率、热力学能；其衍生量可以给出红外光谱、拉曼光谱。它在谱学与热化学上有重要的应用价值，可以服务于药物设计、材料设计等上位学科任务。我们能表明 REST 程序在 RI-JK 近似下的 Skeleton 导数达到相当高的、达到甚至超过其他类似电子结构程序的性能，在 CP-HF 求解过程的 FLOPS 利用率到 40% 以上，同时在算法实现与程序优化上有出色的表现[^1]。而这些是基于 RSTSR 的高效张量计算实现的。

[^1]: Zhu, A. J.; Zhang, I. Y. Efficient Implementation of Hartree-Fock Hessian Skeleton Derivative with Resolution-of-Identity Approximation in REST. International Conference on Theoretical and High Performance Computational Chemistry (ICT-HPCC) 2026.