Merge branch 'devel' into codex/add-yaml-input-support-and-unit-test

caic99 · web-flow · commit 09b3275f9370 · 2025-08-18T10:21:17.000+08:00
diff --git a/.devcontainer/build_cxx.sh b/.devcontainer/build_cxx.sh
@@ -5,7 +5,7 @@ NPROC=$(nproc --all)
 SCRIPT_PATH=$(dirname $(realpath -s $0))
 
 export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
-TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
 
 mkdir -p ${SCRIPT_PATH}/../buildcxx/
 cd ${SCRIPT_PATH}/../buildcxx/
diff --git a/.devcontainer/gdb_lmp b/.devcontainer/gdb_lmp
@@ -2,7 +2,7 @@
 SCRIPT_PATH=$(dirname $(realpath -s $0))
 
 export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
-TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
 
 env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
 	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
diff --git a/.devcontainer/gdb_pytest_lmp b/.devcontainer/gdb_pytest_lmp
@@ -2,7 +2,7 @@
 SCRIPT_PATH=$(dirname $(realpath -s $0))/../..
 
 export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
-TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
 
 env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
 	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
diff --git a/.devcontainer/lmp b/.devcontainer/lmp
@@ -2,7 +2,7 @@
 SCRIPT_PATH=$(dirname $(realpath -s $0))
 
 export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
-TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
 
 env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
 	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
diff --git a/.devcontainer/pytest_lmp b/.devcontainer/pytest_lmp
@@ -2,7 +2,7 @@
 SCRIPT_PATH=$(dirname $(realpath -s $0))/../..
 
 export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch
-TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
 
 env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \
 	LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \
diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml
@@ -26,7 +26,7 @@ jobs:
     - name: Install Python dependencies
       run: |
         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)')
+        export TENSORFLOW_ROOT=$(python -c 'import importlib.util,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 mpich
         source/install/uv_with_retry.sh pip install --system 'torch==2.7' --index-url https://download.pytorch.org/whl/cpu
     - name: Convert models
diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
@@ -46,7 +46,7 @@ jobs:
     - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.7.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)')
+        export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         pip install "paddlepaddle-gpu==3.0.0" -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
         source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
       env:
diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
@@ -27,7 +27,7 @@ jobs:
     - run: |
         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 TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         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 -U setuptools
diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import os
+import re
 import site
 from functools import (
     lru_cache,
@@ -56,6 +57,10 @@ def find_tensorflow() -> tuple[Optional[str], list[str]]:
     ) is not None:
         site_packages = Path(os.environ.get("TENSORFLOW_ROOT")).parent.absolute()
         tf_spec = FileFinder(str(site_packages)).find_spec("tensorflow")
+        if tf_spec is None:
+            raise RuntimeError(
+                f"cannot find TensorFlow under TENSORFLOW_ROOT {os.environ.get('TENSORFLOW_ROOT')}"
+            )
 
     # get tensorflow spec
     # note: isolated build will not work for backend
@@ -153,7 +158,8 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
                 # https://github.com/tensorflow/tensorflow/issues/61830
-                "tensorflow-cpu!=2.15.*; platform_system=='Windows'",
+                # Since TF 2.20, not all symbols are exported to the public API.
+                "tensorflow-cpu!=2.15.*,<2.20; platform_system=='Windows'",
                 # https://github.com/h5py/h5py/issues/2408
                 "h5py>=3.6.0,!=3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
                 *extra_requires,
@@ -228,6 +234,22 @@ def get_tf_version(tf_path: Optional[Union[str, Path]]) -> str:
                 patch = line.split()[-1]
             elif line.startswith("#define TF_VERSION_SUFFIX"):
                 suffix = line.split()[-1].strip('"')
+    if None in (major, minor, patch):
+        # since TF 2.20.0, version information is no more contained in version.h
+        # try to read version from tools/pip_package/setup.py
+        # _VERSION = '2.20.0'
+        setup_file = Path(tf_path) / "tools" / "pip_package" / "setup.py"
+        if setup_file.exists():
+            with open(setup_file) as f:
+                for line in f:
+                    # parse with regex
+                    match = re.search(
+                        r"_VERSION[ \t]*=[ \t]*'(\d+)\.(\d+)\.(\d+)([a-zA-Z0-9]*)?'",
+                        line,
+                    )
+                    if match:
+                        major, minor, patch, suffix = match.groups()
+                        break
     if None in (major, minor, patch):
         raise RuntimeError("Failed to read TF version")
     return ".".join((major, minor, patch)) + suffix
diff --git a/deepmd/pd/model/descriptor/repflows.py b/deepmd/pd/model/descriptor/repflows.py
@@ -69,8 +69,9 @@ def border_op(
         argument8,
     ) -> paddle.Tensor:
         raise NotImplementedError(
-            "border_op is not available since customized Paddle OP library is not built when freezing the model. "
-            "See documentation for DPA3 for details."
+            "The 'border_op' operator is unavailable because the custom Paddle OP library was not built when freezing the model.\n"
+            "To install 'border_op', run: python source/op/pd/setup.py install\n"
+            "For more information, please refer to the DPA3 documentation."
         )
 
     # Note: this hack cannot actually save a model that can be run using LAMMPS.
diff --git a/deepmd/pd/model/descriptor/repformers.py b/deepmd/pd/model/descriptor/repformers.py
@@ -66,8 +66,9 @@ def border_op(
         argument8,
     ) -> paddle.Tensor:
         raise NotImplementedError(
-            "border_op is not available since customized Paddle OP library is not built when freezing the model. "
-            "See documentation for DPA3 for details."
+            "The 'border_op' operator is unavailable because the custom Paddle OP library was not built when freezing the model.\n"
+            "To install 'border_op', run: python source/op/pd/setup.py install\n"
+            "For more information, please refer to the DPA3 documentation."
         )
 
     # Note: this hack cannot actually save a model that can be run using LAMMPS.
diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md
@@ -185,7 +185,10 @@ Switch to the TensorFlow {{ tensorflow_icon }} tab for more information.
 ::::{tab-item} CUDA 12.6
 
 ```bash
-pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
+# release version
+pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
+# nightly-build version
+# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
 pip install deepmd-kit
 ```
 
@@ -194,7 +197,10 @@ pip install deepmd-kit
 ::::{tab-item} CUDA 11.8
 
 ```bash
-pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
+# release version
+pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
+# nightly-build version
+# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu118/
 pip install deepmd-kit
 ```
 
@@ -203,7 +209,10 @@ pip install deepmd-kit
 ::::{tab-item} CPU
 
 ```bash
-pip install paddlepaddle==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
+# release version
+pip install paddlepaddle==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
+# nightly-build version
+# pip install --pre paddlepaddle -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/
 pip install deepmd-kit
 ```
 
diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md
@@ -99,11 +99,22 @@ To install Paddle, run
 
 ```sh
 # cu126
-pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
+# release version
+pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
+# nightly-build version
+# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/
+
 # cu118
-pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
+# release version
+pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
+# nightly-build version
+# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu118/
+
 # cpu
-pip install paddlepaddle==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
+# release version
+pip install paddlepaddle==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
+# nightly-build version
+# pip install --pre paddlepaddle -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/
 ```
 
 :::
diff --git a/doc/model/dpa3.md b/doc/model/dpa3.md
@@ -1,4 +1,4 @@
-# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}
+# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }}
 
 :::{note}
 **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }}
@@ -40,7 +40,11 @@ Virial RMSEs were averaged exclusively for systems containing virial labels (`Al
 
 Note that we set `float32` in all DPA3 models, while `float64` in other models by default.
 
-## Requirements of installation from source code {{ pytorch_icon }}
+## Requirements of installation from source code {{ pytorch_icon }} {{ paddle_icon }}
+
+::::{tab-set}
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
 
 To run the DPA3 model on LAMMPS via source code installation
 (users can skip this step if using [easy installation](../install/easy-install.md)),
@@ -53,6 +57,25 @@ If one runs LAMMPS with MPI, the customized OP library for the C++ interface sho
 If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA,
 otherwise the communication between GPU cards falls back to the slower CPU implementation.
 
+:::
+
+:::{tab-item} Paddle {{ paddle_icon }}
+
+The customized OP library for the Python interface can be installed by
+
+```sh
+cd deepmd-kit/source/op/pd
+python setup.py install
+```
+
+If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI.
+If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA,
+otherwise the communication between GPU cards falls back to the slower CPU implementation.
+
+:::
+
+::::
+
 ## Limitations of the JAX backend with LAMMPS {{ jax_icon }}
 
 When using the JAX backend, 2 or more MPI ranks are not supported. One must set `map` to `yes` using the [`atom_modify`](https://docs.lammps.org/atom_modify.html) command.
diff --git a/doc/train/training.md b/doc/train/training.md
@@ -29,10 +29,10 @@ $ dp --pt train input.json
 :::{tab-item} Paddle {{ paddle_icon }}
 
 ```bash
-# training model in eager mode
+# training model
 $ dp --pd train input.json
 
-# [experimental] training model with CINN compiler for better performance,
+# [experimental] training models with the CINN compiler (~40%+ speedup)
 # see: https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/paddle_v3_features/cinn_cn.html
 ## If the shape(s) of batch input data are dynamic during training(default).
 $ CINN=1 dp --pd train input.json
diff --git a/source/cmake/Findtensorflow.cmake b/source/cmake/Findtensorflow.cmake
@@ -291,8 +291,10 @@ if(NOT DEFINED TENSORFLOW_VERSION)
     TENSORFLOW_VERSION_RUN_RESULT_VAR TENSORFLOW_VERSION_COMPILE_RESULT_VAR
     ${CMAKE_CURRENT_BINARY_DIR}/tf_version
     "${CMAKE_CURRENT_LIST_DIR}/tf_version.cpp"
-    CMAKE_FLAGS "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}"
-    RUN_OUTPUT_VARIABLE TENSORFLOW_VERSION
+    CMAKE_FLAGS
+      "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}" LINK_LIBRARIES
+      ${TensorFlowFramework_LIBRARY} ${TensorFlow_LIBRARY}
+      RUN_OUTPUT_STDOUT_VARIABLE TENSORFLOW_VERSION
     COMPILE_OUTPUT_VARIABLE TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR)
   if(NOT ${TENSORFLOW_VERSION_COMPILE_RESULT_VAR})
     message(
@@ -304,6 +306,23 @@ if(NOT DEFINED TENSORFLOW_VERSION)
   endif()
 endif()
 
+if(TENSORFLOW_VERSION VERSION_GREATER_EQUAL 2.20)
+  # since TF 2.20, macros like TF_MAJOR_VERSION, TF_MINOR_VERSION, and
+  # TF_PATCH_VERSION are not defined We manuanlly define them in our CMake files
+  # first, split TENSORFLOW_VERSION (e.g. 2.20.0rc0) to 2 20 0 rc0
+  string(REGEX MATCH "^([0-9]+)\\.([0-9]+)\\.([0-9]+)(.*)$" _match
+               ${TENSORFLOW_VERSION})
+  if(_match)
+    set(TF_MAJOR_VERSION ${CMAKE_MATCH_1})
+    set(TF_MINOR_VERSION ${CMAKE_MATCH_2})
+    set(TF_PATCH_VERSION ${CMAKE_MATCH_3})
+    # add defines
+    add_definitions(-DTF_MAJOR_VERSION=${TF_MAJOR_VERSION})
+    add_definitions(-DTF_MINOR_VERSION=${TF_MINOR_VERSION})
+    add_definitions(-DTF_PATCH_VERSION=${TF_PATCH_VERSION})
+  endif()
+endif()
+
 # print message
 if(NOT TensorFlow_FIND_QUIETLY)
   message(
diff --git a/source/cmake/tf_version.cpp b/source/cmake/tf_version.cpp
@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <iostream>
 
-#include "tensorflow/core/public/version.h"
+#include "tensorflow/c/c_api.h"
 
 int main(int argc, char* argv[]) {
   // See
   // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/public/version.h
   // TF_VERSION_STRING has been available since TensorFlow v0.6
-  std::cout << TF_VERSION_STRING;
+  // Aug 2025: since TF 2.20, TF_VERSION_STRING is no more available;
+  // try to use the C API TF_Version
+  std::cout << TF_Version();
   return 0;
 }
diff --git a/source/tests/pt/test_padding_atoms.py b/source/tests/pt/test_padding_atoms.py
@@ -49,7 +49,8 @@ def setUp(self) -> None:
         self.sel = [16, 8]
         self.rcut = 2.2
         self.rcut_smth = 0.4
-        self.atol = 1e-8
+        self.atol = 1e-6
+        self.rtol = 1e-5
 
 
 class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist):
@@ -79,6 +80,7 @@ def test_padding_atoms_consistency(self):
             to_numpy_array(result[var_name].cpu().detach()),
             np.mean(to_numpy_array(result[f"atom_{var_name}"].cpu().detach()), axis=1),
             atol=self.atol,
+            rtol=self.rtol,
         )
         # test padding atoms
         padding_atoms_list = [1, 5, 10]
@@ -105,6 +107,7 @@ def test_padding_atoms_consistency(self):
                 to_numpy_array(result[var_name].cpu().detach()),
                 to_numpy_array(result_padding[var_name].cpu().detach()),
                 atol=self.atol,
+                rtol=self.rtol,
             )
 
 
