lrux: 在JAX中实现快速低秩更新行列式和佩尔逊行列式
lrux: Fast low-rank updates of determinants and Pfaffians in JAX
AI总结:
lrux通过低秩更新技术在JAX中实现行列式和佩尔逊行列式的高效计算,显著提升量子蒙特卡洛算法的性能。
摘要:
我们介绍了lrux,一个基于JAX的软件包,用于快速低秩更新行列式和佩尔逊行列式,旨在解决各种量子蒙特卡洛(QMC)算法中的主要计算瓶颈。该软件包实现了高效的低秩更新,将连续波函数评估的成本从O(n³)降低到O(n²k),当更新秩k小于矩阵维度n时。支持行列式和佩尔逊行列式的更新,同时支持延迟更新策略,通过交换浮点运算来减少现代加速器上的内存交通。lrux原生集成了JAX转换,如JIT编译、向量化和自动微分,并支持实数和复数数据类型。在GPU上的基准测试显示,在大矩阵尺寸下可达到高达1000倍的加速。lrux使抗对称波函数的可扩展、高性能评估成为可能,并设计为QMC工作流程中的即插即用组件。lrux可在https://github.com/ChenAo-Phys/lrux上获得。
英文摘要:
We present lrux, a JAX-based software package for fast low-rank updates of determinants and Pfaffians, targeting the dominant computational bottleneck in various quantum Monte Carlo (QMC) algorithms. The package implements efficient low-rank updates that reduce the cost of successive wavefunction evaluations from $\mathcal{O}(n^3)$ to $\mathcal{O}(n^2k)$ when the update rank $k$ is smaller than the dimension $n$ of matrices. Both determinant and Pfaffian updates are supported, together with delayed-update strategies that trade floating-point operations for reduced memory traffic on modern accelerators. lrux natively integrates with JAX transformations such as JIT compilation, vectorization, and automatic differentiation, and supports both real and complex data types. Benchmarks on GPUs demonstrate up to $1000\times$ speedup at large matrix sizes. lrux enables scalable, high-performance evaluation of antisymmetric wavefunctions and is designed as a drop-in component for a wide range of QMC workflows. lrux is available at https://github.com/ChenAo-Phys/lrux.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。