chore: merge master into fix-paddle-download-hash

Bring the pull request up to date with the current base branch.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh

Author: njzjz-bot

.github/workflows/test_cuda.yml

@@ -19,29 +19,13 @@ jobs:
     runs-on: gpu
     # The serial Python test suite can take several hours on the GPU runner.
     timeout-minutes: 480
-    # Share ONE AOTInductor on-disk compile cache across the two pytest lanes;
-    # actions/cache (below) persists it across runs so unchanged models never
-    # recompile.
-    env:
-      TORCHINDUCTOR_CACHE_DIR: ${{ github.workspace }}/.inductor-cache
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     # container:
     #   image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
     #   options: --gpus all
     if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
       - uses: actions/checkout@v7
-      - name: Cache AOTInductor compile artifacts
-        # Restored AFTER checkout so checkout's git-clean can't wipe it. The repo
-        # Actions cache is capped at 10 GB; if the inductor cache outgrows it the
-        # oldest entries are evicted -- correctness is unaffected (inductor
-        # re-validates each graph; a miss just recompiles).
-        uses: actions/cache@v6
-        with:
-          path: ${{ github.workspace }}/.inductor-cache
-          key: inductor-cuda-${{ github.sha }}
-          restore-keys: |
-            inductor-cuda-
       - uses: actions/setup-python@v7
         with:
           python-version: "3.11"
@@ -78,84 +62,23 @@ jobs:
           DP_ENABLE_NATIVE_OPTIMIZATION: 1
           DP_ENABLE_PYTORCH: 1
       - run: dp --version
-      - name: Run Python tests (overlap CPU-bound aoti compiles with the GPU lane)
-        # Two pytest lanes on one GPU box:
-        #   Lane 1 (CPU stream): the .pt2-freezing tests (-m aoti_compile) --
-        #     CPU-bound, GPU ~98% idle while inductor/g++/ptxas run.
-        #   Lane 2 (GPU stream): the GPU-bound remainder (-m "not aoti_compile")
-        #     -- the bulk of the suite and the wall-clock bottleneck. It does not
-        #     freeze .pt2, so it only reads the inductor cache, never competing
-        #     to write.
-        # Wall-clock becomes max(lane 1, lane 2) instead of their sum; the whole
-        # compile stream is hidden behind the GPU lane. (The C++ build + fixture
-        # generation runs in the separate test_cc job.)
-        run: |
-          # Cap the compile lane's inductor threads to ~half the cores so lane 2
-          # keeps CPU for its host-side GPU dispatch; nice the compile work too.
-          nc=$(( $(nproc) / 2 )); [ "$nc" -lt 1 ] && nc=1
-          # rc=0 + `|| rc=$?` so a lane's failure does not let the step's default
-          # `set -e` skip the wall-clock echo (the timing we most want on failure).
-          ( ta=$(date +%s); rc=0
-            TORCHINDUCTOR_COMPILE_THREADS=$nc nice -n 15 \
-              python -m pytest source/tests -m "aoti_compile" \
-                -p no:cacheprovider --junit-xml=junit-aoti.xml || rc=$?
-            # pytest exit 5 == "no tests collected"; harmless for the compile lane.
-            [ "$rc" = 5 ] && rc=0
-            echo "[lane aoti_compile] wall=$(( $(date +%s) - ta ))s exit=$rc"
-            exit $rc ) &
-          pid_a=$!
-          ( tb=$(date +%s); rc=0
-            python -m pytest source/tests -m "not aoti_compile" \
-              -p no:cacheprovider --junit-xml=junit-gpu.xml || rc=$?
-            echo "[lane gpu] wall=$(( $(date +%s) - tb ))s exit=$rc"
-            exit $rc ) &
-          pid_b=$!
-          rc_a=0; wait $pid_a || rc_a=$?
-          rc_b=0; wait $pid_b || rc_b=$?
-          echo "lane exit codes: aoti_compile=$rc_a gpu=$rc_b"
-          exit $(( rc_a != 0 || rc_b != 0 ))
+      - run: python -m pytest source/tests
         env:
           NUM_WORKERS: 0
           CUDA_VISIBLE_DEVICES: 0
           # See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
           XLA_PYTHON_CLIENT_PREALLOCATE: false
           XLA_PYTHON_CLIENT_ALLOCATOR: platform
           FLAGS_use_stride_compute_kernel: 0
-      - name: Upload per-lane JUnit reports
-        # junit-aoti.xml = compile lane (T_A), junit-gpu.xml = GPU lane (T_B);
-        # used to compare the two lanes against the pre-split serial baseline.
-        if: always()
-        uses: actions/upload-artifact@v7
-        with:
-          name: cuda-pytest-junit
-          path: |
-            junit-aoti.xml
-            junit-gpu.xml
-          if-no-files-found: warn
 
   test_cc:
     name: Test C++ on CUDA
     runs-on: gpu
-    # The gen_*.py fixture builds drive the same AOTInductor backend on the same
-    # architectures as the Python lanes, so they reuse the kernels cached by
-    # test_python / previous runs.
-    env:
-      TORCHINDUCTOR_CACHE_DIR: ${{ github.workspace }}/.inductor-cache
     if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
       # Jobs run on separate runners, so the C++ job needs its own complete
       # CUDA toolchain and Python dependency installation.
       - uses: actions/checkout@v7
-      - name: Cache AOTInductor compile artifacts
-        # Restored AFTER checkout so checkout's git-clean can't wipe it. Both
-        # CUDA jobs share the same key; if both try to save it, the second save
-        # is skipped with a warning (caches are immutable), which is harmless.
-        uses: actions/cache@v6
-        with:
-          path: ${{ github.workspace }}/.inductor-cache
-          key: inductor-cuda-${{ github.sha }}
-          restore-keys: |
-            inductor-cuda-
       - uses: actions/setup-python@v7
         with:
           python-version: "3.11"

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.22
+    rev: v0.16.0
     hooks:
       - id: ruff
         args: ["--fix"]

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -758,7 +758,7 @@ def change_out_bias(
             delta_bias, out_std = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 model_forward=self._get_forward_wrapper_func(),
                 rcond=self.rcond,
@@ -771,7 +771,7 @@ def change_out_bias(
             bias_out, std_out = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,

deepmd/dpmodel/descriptor/descriptor.py

@@ -110,6 +110,37 @@ def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""
         raise NotImplementedError
 
+    def set_stat_mean_and_stddev(
+        self,
+        mean: Array,
+        stddev: Array,
+    ) -> None:
+        """Update the normalization arrays of the descriptor block.
+
+        Parameters
+        ----------
+        mean
+            Mean of the environment matrix.
+        stddev
+            Standard deviation of the environment matrix.
+
+        Returns
+        -------
+        None
+        """
+        self["davg"] = mean
+        self["dstd"] = stddev
+
+    def get_stat_mean_and_stddev(self) -> tuple[Array, Array]:
+        """Return the normalization arrays of the descriptor block.
+
+        Returns
+        -------
+        tuple[Array, Array]
+            Mean and standard deviation of the environment matrix.
+        """
+        return self["davg"], self["dstd"]
+
     def share_params(
         self, base_class: Any, shared_level: Any, resume: bool = False
     ) -> None:

deepmd/dpmodel/descriptor/dpa1.py

@@ -1461,15 +1461,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self, use_graph=True)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/repflows.py

@@ -494,15 +494,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/repformers.py

@@ -459,15 +459,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.stddev)

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -349,15 +349,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)

deepmd/dpmodel/descriptor/se_r.py

@@ -328,15 +328,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)

deepmd/dpmodel/descriptor/se_t.py

@@ -303,15 +303,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         xp = array_api_compat.array_namespace(self.dstd)