[r3.0] release v3.0.2 (#4589)

This patch version only contains bug fixes, enhancements, and
documentation improvements.



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
	- Added Sphinx-powered documentation support.
	- Introduced a new test suite for dipole model evaluation.

- **Bug Fixes**
	- Enhanced error handling for GPU support in custom operations.

- **Refactor**
- Upgraded key dependencies to version 2.6.0 and refined tensor
operations for improved clarity and performance.
- Streamlined continuous integration settings, including simplified
runner selection for ARM builds and enhanced error handling.

- **Documentation**
- Updated default environment settings and reference links to reflect
current publications.
	- Added a new section on using the PyTorch Profiler with TensorBoard.

- **Tests & Chores**
- Refined test comparisons and adjusted build configurations for
enhanced stability.
	- Added new import for `torch._dynamo` in test files.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: njzjz

.devcontainer/download_libtorch.sh

@@ -4,5 +4,5 @@ set -ev
 SCRIPT_PATH=$(dirname $(realpath -s $0))
 cd ${SCRIPT_PATH}/..
 
-wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcpu.zip -O ~/libtorch.zip
+wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcpu.zip -O ~/libtorch.zip
 unzip ~/libtorch.zip

.github/workflows/build_wheel.yml

@@ -14,24 +14,8 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  determine-arm64-runner:
-    runs-on: ubuntu-latest
-    permissions: read-all
-    outputs:
-      runner: ${{ steps.set-runner.outputs.runner }}
-    steps:
-      - name: Determine which runner to use for ARM64 build
-        id: set-runner
-        run: |
-          if [ "${{ github.repository_owner }}" == "deepmodeling" ]; then
-            echo "runner=[\"Linux\",\"ARM64\"]" >> $GITHUB_OUTPUT
-          else
-            echo "runner=\"ubuntu-latest\"" >> $GITHUB_OUTPUT
-          fi
-
   build_wheels:
     name: Build wheels for cp${{ matrix.python }}-${{ matrix.platform_id }}
-    needs: determine-arm64-runner
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -65,7 +49,7 @@ jobs:
             platform_id: win_amd64
             dp_variant: cpu
           # linux-aarch64
-          - os: ${{ fromJson(needs.determine-arm64-runner.outputs.runner) }}
+          - os: ubuntu-24.04-arm
             python: 310
             platform_id: manylinux_aarch64
             dp_variant: cpu

.github/workflows/test_cc.yml

@@ -29,7 +29,7 @@ jobs:
     - run: python -m pip install uv
     - name: Install Python dependencies
       run: |
-        source/install/uv_with_retry.sh pip install --system tensorflow-cpu
+        source/install/uv_with_retry.sh pip install --system tensorflow-cpu~=2.18.0 jax==0.5.0
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py
     - name: Convert models

.github/workflows/test_cuda.yml

@@ -47,11 +47,11 @@ jobs:
          && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
       if: false  # skip as we use nvidia image
     - run: python -m pip install -U uv
-    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.5.0" "jax[cuda12]"
+    - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.6.0" "jax[cuda12]==0.5.0"
     - run: |
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-        source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py
+        source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
@@ -67,7 +67,7 @@ jobs:
       run: source/tests/infer/convert-models.sh
     - name: Download libtorch
       run: |
-         wget https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcu124.zip -O libtorch.zip
+         wget https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu124.zip -O libtorch.zip
          unzip libtorch.zip
     - run: |
         export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch

.github/workflows/test_python.yml

@@ -25,11 +25,11 @@ jobs:
         python-version: ${{ matrix.python }}
     - run: python -m pip install -U uv
     - run: |
-        source/install/uv_with_retry.sh pip install --system openmpi tensorflow-cpu
+        source/install/uv_with_retry.sh pip install --system openmpi tensorflow-cpu~=2.18.0
         source/install/uv_with_retry.sh pip install --system torch -i https://download.pytorch.org/whl/cpu
         export TENSORFLOW_ROOT=$(python -c 'import tensorflow;print(tensorflow.__path__[0])')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py
+        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.5.0;python_version>='3.10'"
         source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
       env:
         # Please note that uv has some issues with finding

.readthedocs.yml

@@ -10,5 +10,7 @@ build:
       - VIRTUAL_ENV=$READTHEDOCS_VIRTUALENV_PATH uv pip install -r doc/requirements.txt
   apt_packages:
     - inkscape
+sphinx:
+  configuration: doc/conf.py
 formats:
   - pdf

CITATIONS.bib

@@ -40,6 +40,31 @@ @article{Zeng_JChemPhys_2023_v159_p054801
   doi          = {10.1063/5.0155600},
 }
 
+@article{Zeng_arXiv_2025_p2502.19161,
+  annote       = {general purpose},
+  author       = {
+    Jinzhe Zeng and Duo Zhang and Anyang Peng and Xiangyu Zhang and Sensen He
+    and Yan Wang and Xinzijian Liu and Hangrui Bi and Yifan Li and Chun Cai and
+    Chengqian Zhang and Yiming Du and Jia-Xin Zhu and Pinghui Mo and Zhengtao
+    Huang and Qiyu Zeng and Shaochen Shi and Xuejian Qin and Zhaoxi Yu and
+    Chenxing Luo and Ye Ding and Yun-Pei Liu and Ruosong Shi and Zhenyu Wang
+    and Sigbj{\o}rn L{\o}land Bore and Junhan Chang and Zhe Deng and Zhaohan
+    Ding and Siyuan Han and Wanrun Jiang and Guolin Ke and Zhaoqing Liu and
+    Denghui Lu and Koki Muraoka and Hananeh Oliaei and Anurag Kumar Singh and
+    Haohui Que and Weihong Xu and Zhangmancang Xu and Yong-Bin Zhuang and Jiayu
+    Dai and Timothy J. Giese and Weile Jia and Ben Xu and Darrin M. York and
+    Linfeng Zhang and Han Wang
+  },
+  title        = {
+    {DeePMD-kit v3: A Multiple-Backend Framework for Machine Learning
+    Potentials}
+  },
+  journal      = {arXiv},
+  year         = 2025,
+  pages        = {2502.19161},
+  doi          = {10.48550/arXiv.2502.19161},
+}
+
 @article{Lu_CompPhysCommun_2021_v259_p107624,
   annote       = {GPU support},
   title        = {

README.md

@@ -36,6 +36,7 @@ If you use this code in any future publications, please cite the following publi
 - Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang. "DeePMD-kit v2: A software package for deep potential models." J. Chem. Phys. 159 (2023): 054801.
   [![doi:10.1063/5.0155600](https://img.shields.io/badge/DOI-10.1063%2F5.0155600-blue)](https://doi.org/10.1063/5.0155600)
   [![Citations](https://citations.njzjz.win/10.1063/5.0155600)](https://badge.dimensions.ai/details/doi/10.1063/5.0155600)
+- Jinzhe Zeng, Duo Zhang, Anyang Peng, Xiangyu Zhang, Sensen He, Yan Wang, Xinzijian Liu, Hangrui Bi, Yifan Li, Chun Cai, Chengqian Zhang, Yiming Du, Jia-Xin Zhu, Pinghui Mo, Zhengtao Huang, Qiyu Zeng, Shaochen Shi, Xuejian Qin, Zhaoxi Yu, Chenxing Luo, Ye Ding, Yun-Pei Liu, Ruosong Shi, Zhenyu Wang, Sigbjørn Løland Bore, Junhan Chang, Zhe Deng, Zhaohan Ding, Siyuan Han, Wanrun Jiang, Guolin Ke, Zhaoqing Liu, Denghui Lu, Koki Muraoka, Hananeh Oliaei, Anurag Kumar Singh, Haohui Que, Weihong Xu, Zhangmancang Xu, Yong-Bin Zhuang, Jiayu Dai, Timothy J. Giese, Weile Jia, Ben Xu, Darrin M. York, Linfeng Zhang, Han Wang. "DeePMD-kit v3: A Multiple-Backend Framework for Machine Learning Potentials." [arXiv:2502.19161](https://arxiv.org/abs/2502.19161).
 
 In addition, please follow [the bib file](CITATIONS.bib) to cite the methods you used.
 
@@ -68,14 +69,16 @@ In addition to building up potential energy models, DeePMD-kit can also be used
 - Non-von-Neumann.
 - C API to interface with the third-party packages.
 
-See [our latest paper](https://doi.org/10.1063/5.0155600) for details of all features until v2.2.3.
+See [our v2 paper](https://doi.org/10.1063/5.0155600) for details of all features until v2.2.3.
 
 #### v3
 
 - Multiple backends supported. Add PyTorch and JAX backends.
 - The DPA-2 model.
 - Plugin mechanisms for external models.
 
+See [our v3 paper](https://doi.org/10.48550/arXiv.2502.19161) for details of all features until v3.0.
+
 ## Install and use DeePMD-kit
 
 Please read the [online documentation](https://deepmd.readthedocs.io/) for how to install and use DeePMD-kit.

backend/find_pytorch.py

@@ -116,7 +116,7 @@ def get_pt_requirement(pt_version: str = "") -> dict:
         cuda_version = os.environ.get("CUDA_VERSION", "12.2")
         if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
             # CUDA 12.2, cudnn 9
-            pt_version = "2.5.0"
+            pt_version = "2.6.0"
         elif cuda_version in SpecifierSet(">=11,<12"):
             # CUDA 11.8, cudnn 8
             pt_version = "2.3.1"

deepmd/dpmodel/descriptor/dpa1.py

@@ -899,6 +899,7 @@ def call(
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # nfnl x nnei
         exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
+        exclude_mask = xp.astype(exclude_mask, xp.bool)
         # nfnl x nnei
         nlist = xp.reshape(nlist, (nf * nloc, nnei))
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))