Merge branch 'devel' into spin_lmp

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

.github/workflows/build_cc.yml

@@ -11,7 +11,7 @@ name: Build C++
 jobs:
   buildcc:
     name: Build C++
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         include:

.github/workflows/codeql.yml

@@ -13,7 +13,7 @@ concurrency:
 jobs:
   analyze:
     name: Analyze
-    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }}
+    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-22.04' }}
     timeout-minutes: ${{ (matrix.language == 'swift' && 120) || 360 }}
     permissions:
       actions: read

.github/workflows/test_cc.yml

@@ -11,7 +11,7 @@ name: Test C++
 jobs:
   testcc:
     name: Test C++
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         check_memleak: [true, false]

.github/workflows/test_cuda.yml

@@ -51,7 +51,7 @@ jobs:
     - 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] mpi4py
+        source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1

.github/workflows/test_python.yml

@@ -16,7 +16,7 @@ jobs:
       fail-fast: false
       matrix:
         group: [1, 2, 3, 4, 5, 6]
-        python: ["3.8", "3.12"]
+        python: ["3.9", "3.12"]
 
     steps:
     - uses: actions/checkout@v4
@@ -28,13 +28,13 @@ jobs:
         source/install/uv_with_retry.sh pip install --system mpich
         source/install/uv_with_retry.sh pip install --system "torch==2.3.0+cpu.cxx11.abi" -i https://download.pytorch.org/whl/
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-        source/install/uv_with_retry.sh pip install --system --only-binary=horovod -e .[cpu,test] horovod[tensorflow-cpu] mpi4py
+        source/install/uv_with_retry.sh pip install --system --only-binary=horovod -e .[cpu,test,jax] horovod[tensorflow-cpu] mpi4py
       env:
         # Please note that uv has some issues with finding
         # existing TensorFlow package. Currently, it uses
         # TensorFlow in the build dependency, but if it
         # changes, setting `TENSORFLOW_ROOT`.
-        TENSORFLOW_VERSION: ${{ matrix.python == '3.8' && '2.13.1' || '2.16.1' }}
+        TENSORFLOW_VERSION: 2.16.1
         DP_ENABLE_PYTORCH: 1
         DP_BUILD_TESTING: 1
         UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/njzjz/simple https://pypi.anaconda.org/mpi4py/simple"
@@ -69,7 +69,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python: ["3.8", "3.12"]
+        python: ["3.9", "3.12"]
     needs: testpython
     steps:
     - name: Get durations from cache

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
         exclude: "^.+\\.pbtxt$"
@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.5
+    rev: v0.6.9
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -47,15 +47,15 @@ repos:
         exclude: ^source/3rdparty
   # Python inside docs
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.18.0
+    rev: 1.19.0
     hooks:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.8
+    rev: v19.1.1
     hooks:
       - id: clang-format
-        exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc
+        exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$)
   # markdown, yaml, CSS, javascript
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v4.0.0-alpha.8
@@ -146,7 +146,7 @@ repos:
         exclude: .pre-commit-config.yaml|source/lmp
   # customized pylint rules
   - repo: https://github.com/pylint-dev/pylint/
-    rev: v3.2.7
+    rev: v3.3.1
     hooks:
       - id: pylint
         entry: env PYTHONPATH=source/checker pylint

backend/dp_backend.py

@@ -1,10 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """A PEP-517 backend to find TensorFlow."""
 
-from typing import (
-    List,
-)
-
 from scikit_build_core import build as _orig
 
 from .find_pytorch import (
@@ -26,7 +22,7 @@
 ]
 
 
-def __dir__() -> List[str]:
+def __dir__() -> list[str]:
     return __all__
 
 
@@ -42,7 +38,7 @@ def __dir__() -> List[str]:
 
 def get_requires_for_build_wheel(
     config_settings: dict,
-) -> List[str]:
+) -> list[str]:
     return (
         _orig.get_requires_for_build_wheel(config_settings)
         + find_tensorflow()[1]
@@ -52,7 +48,7 @@ def get_requires_for_build_wheel(
 
 def get_requires_for_build_editable(
     config_settings: dict,
-) -> List[str]:
+) -> list[str]:
     return (
         _orig.get_requires_for_build_editable(config_settings)
         + find_tensorflow()[1]

backend/dynamic_metadata.py

@@ -4,8 +4,6 @@
     Path,
 )
 from typing import (
-    Dict,
-    List,
     Optional,
 )
 
@@ -27,13 +25,13 @@
 __all__ = ["dynamic_metadata"]
 
 
-def __dir__() -> List[str]:
+def __dir__() -> list[str]:
     return __all__
 
 
 def dynamic_metadata(
     field: str,
-    settings: Optional[Dict[str, object]] = None,
+    settings: Optional[dict[str, object]] = None,
 ):
     assert field in ["optional-dependencies", "entry-points", "scripts"]
     _, _, find_libpython_requires, extra_scripts, tf_version, pt_version = (

backend/find_pytorch.py

@@ -18,9 +18,7 @@
     get_path,
 )
 from typing import (
-    List,
     Optional,
-    Tuple,
     Union,
 )
 
@@ -30,7 +28,7 @@
 
 
 @lru_cache
-def find_pytorch() -> Tuple[Optional[str], List[str]]:
+def find_pytorch() -> tuple[Optional[str], list[str]]:
     """Find PyTorch library.
 
     Tries to find PyTorch in the order of:

backend/find_tensorflow.py

@@ -17,9 +17,7 @@
     get_path,
 )
 from typing import (
-    List,
     Optional,
-    Tuple,
     Union,
 )
 
@@ -29,7 +27,7 @@
 
 
 @lru_cache
-def find_tensorflow() -> Tuple[Optional[str], List[str]]:
+def find_tensorflow() -> tuple[Optional[str], list[str]]:
     """Find TensorFlow library.
 
     Tries to find TensorFlow in the order of:
@@ -156,18 +154,15 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "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'",
-                # TODO: build(wheel): unpin h5py on aarch64
-                # Revert after https://github.com/h5py/h5py/issues/2408 is fixed;
-                # or set UV_PREFER_BINARY when https://github.com/astral-sh/uv/issues/1794 is resolved.
-                # 3.6.0 is the first version to have aarch64 wheels.
-                "h5py>=3.6.0,<3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
+                # https://github.com/h5py/h5py/issues/2408
+                "h5py>=3.6.0,!=3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
                 *extra_requires,
             ],
             "gpu": [
                 "tensorflow",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
                 # See above.
-                "h5py>=3.6.0,<3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
+                "h5py>=3.6.0,!=3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
                 *extra_requires,
             ],
             **extra_select,

backend/read_env.py

@@ -5,9 +5,6 @@
 from functools import (
     lru_cache,
 )
-from typing import (
-    Tuple,
-)
 
 from packaging.version import (
     Version,
@@ -24,7 +21,7 @@
 
 
 @lru_cache
-def get_argument_from_env() -> Tuple[str, list, list, dict, str, str]:
+def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
     """Get the arguments from environment variables.
 
     The environment variables are assumed to be not changed during the build.

deepmd/backend/backend.py

@@ -10,9 +10,6 @@
     TYPE_CHECKING,
     Callable,
     ClassVar,
-    Dict,
-    List,
-    Type,
 )
 
 from deepmd.utils.plugin import (
@@ -45,7 +42,7 @@ class Backend(PluginVariant, make_plugin_registry("backend")):
     """
 
     @staticmethod
-    def get_backend(key: str) -> Type["Backend"]:
+    def get_backend(key: str) -> type["Backend"]:
         """Get the backend by key.
 
         Parameters
@@ -61,7 +58,7 @@ def get_backend(key: str) -> Type["Backend"]:
         return Backend.get_class_by_type(key)
 
     @staticmethod
-    def get_backends() -> Dict[str, Type["Backend"]]:
+    def get_backends() -> dict[str, type["Backend"]]:
         """Get all the registered backend names.
 
         Returns
@@ -74,7 +71,7 @@ def get_backends() -> Dict[str, Type["Backend"]]:
     @staticmethod
     def get_backends_by_feature(
         feature: "Backend.Feature",
-    ) -> Dict[str, Type["Backend"]]:
+    ) -> dict[str, type["Backend"]]:
         """Get all the registered backend names with a specific feature.
 
         Parameters
@@ -94,7 +91,7 @@ def get_backends_by_feature(
         }
 
     @staticmethod
-    def detect_backend_by_model(filename: str) -> Type["Backend"]:
+    def detect_backend_by_model(filename: str) -> type["Backend"]:
         """Detect the backend of the given model file.
 
         Parameters
@@ -128,7 +125,7 @@ class Feature(Flag):
 
     features: ClassVar[Feature] = Feature(0)
     """The features of the backend."""
-    suffixes: ClassVar[List[str]] = []
+    suffixes: ClassVar[list[str]] = []
     """The supported suffixes of the saved model.
 
     The first element is considered as the default suffix."""
@@ -157,7 +154,7 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
 
     @property
     @abstractmethod
-    def deep_eval(self) -> Type["DeepEvalBackend"]:
+    def deep_eval(self) -> type["DeepEvalBackend"]:
         """The Deep Eval backend of the backend.
 
         Returns
@@ -169,7 +166,7 @@ def deep_eval(self) -> Type["DeepEvalBackend"]:
 
     @property
     @abstractmethod
-    def neighbor_stat(self) -> Type["NeighborStat"]:
+    def neighbor_stat(self) -> type["NeighborStat"]:
         """The neighbor statistics of the backend.
 
         Returns

deepmd/backend/dpmodel.py

@@ -3,8 +3,6 @@
     TYPE_CHECKING,
     Callable,
     ClassVar,
-    List,
-    Type,
 )
 
 from deepmd.backend.backend import (
@@ -37,7 +35,7 @@ class DPModelBackend(Backend):
         Backend.Feature.DEEP_EVAL | Backend.Feature.NEIGHBOR_STAT | Backend.Feature.IO
     )
     """The features of the backend."""
-    suffixes: ClassVar[List[str]] = [".dp", ".yaml", ".yml"]
+    suffixes: ClassVar[list[str]] = [".dp", ".yaml", ".yml"]
     """The suffixes of the backend."""
 
     def is_available(self) -> bool:
@@ -62,7 +60,7 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
         raise NotImplementedError(f"Unsupported backend: {self.name}")
 
     @property
-    def deep_eval(self) -> Type["DeepEvalBackend"]:
+    def deep_eval(self) -> type["DeepEvalBackend"]:
         """The Deep Eval backend of the backend.
 
         Returns
@@ -77,7 +75,7 @@ def deep_eval(self) -> Type["DeepEvalBackend"]:
         return DeepEval
 
     @property
-    def neighbor_stat(self) -> Type["NeighborStat"]:
+    def neighbor_stat(self) -> type["NeighborStat"]:
         """The neighbor statistics of the backend.
 
         Returns