feat(pt): Add DPA4/SeZM descriptor & model 🎉🎉🎉

.github/workflows/copilot-setup-steps.yml

@@ -52,7 +52,7 @@ jobs:
         run: uv pip install --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
 
       - name: Build Python package
-        run: uv pip install -e .[cpu,test]
+        run: uv pip install -e .[cpu,test,torch]
 
       - name: Install prek tools
         run: uv tool install prek

.github/workflows/test_cc.yml

@@ -46,7 +46,10 @@ jobs:
         run: |
           source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax_cpu --torch-backend cpu
           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
+          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
+          source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax,torch] mpi4py mpich
+        env:
+          DP_ENABLE_PYTORCH: 1
       - name: Convert models
         run: source/tests/infer/convert-models.sh
       # https://github.com/actions/runner-images/issues/9491

.github/workflows/test_python.yml

@@ -31,7 +31,7 @@ jobs:
           source/install/uv_with_retry.sh pip install --system openmpi --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
           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 --group pin_jax_cpu
+          source/install/uv_with_retry.sh pip install --system -e .[test,jax,torch] mpi4py --group pin_jax_cpu
           source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn paddlepaddle==3.4.0.dev20260310
         env:
           # Please note that uv has some issues with finding

backend/find_pytorch.py

@@ -136,6 +136,7 @@ def get_pt_requirement(pt_version: str = "") -> dict:
             if pt_version != ""
             # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc
             else "torch>=2.1.0",
+            "e3nn>=0.5.9",
             *mpi_requirement,
             *cibw_requirement,
         ],

deepmd/dpmodel/utils/dist_check.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Minimum pairwise distance check for frame validity filtering."""
+
+from __future__ import (
+    annotations,
+)
+
+import numpy as np
+
+_MIN_PAIR_DIST_BLOCK_PAIRS = 262_144
+
+
+def compute_min_pair_dist_single(
+    coord: np.ndarray,
+    box: np.ndarray | None,
+    atype: np.ndarray,
+    stop_below: float | None = None,
+) -> float:
+    """Compute the minimum pairwise atomic distance for a single frame.
+
+    Parameters
+    ----------
+    coord : np.ndarray
+        Atomic coordinates, flattened with shape (natoms * 3,)
+        or reshaped as (natoms, 3).
+    box : np.ndarray or None
+        Box vectors with shape (9,) for PBC, or None for non-PBC.
+    atype : np.ndarray
+        Atom types with shape (natoms,). Virtual atoms (type < 0)
+        are excluded from the distance check.
+    stop_below : float or None
+        Optional early-stop threshold. If a block has any pair closer
+        than this value, the block minimum is returned immediately.
+
+    Returns
+    -------
+    float
+        Minimum pairwise distance. Returns inf if fewer than 2
+        real atoms exist.
+    """
+    coord = coord.reshape(-1, 3)
+
+    # === Step 1. Filter out virtual atoms ===
+    real_mask = atype.ravel() >= 0
+    real_coord = coord[real_mask]
+    n_real = real_coord.shape[0]
+    if n_real < 2:
+        return float("inf")
+
+    # === Step 2. Prepare minimum image convention for PBC ===
+    if box is not None:
+        cell = box.reshape(3, 3)
+        inv_cell = np.linalg.inv(cell)
+    else:
+        cell = None
+        inv_cell = None
+
+    # === Step 3. Compute distances in bounded row blocks ===
+    block_size = max(1, min(n_real, _MIN_PAIR_DIST_BLOCK_PAIRS // n_real))
+    min_dist_sq = float("inf")
+    stop_dist_sq = (
+        float(stop_below) * float(stop_below)
+        if stop_below is not None and stop_below > 0.0
+        else None
+    )
+    for start in range(0, n_real, block_size):
+        stop = min(start + block_size, n_real)
+        diff = real_coord[np.newaxis, :, :] - real_coord[start:stop, np.newaxis, :]
+
+        if cell is not None and inv_cell is not None:
+            frac_diff = diff @ inv_cell
+            frac_diff -= np.round(frac_diff)
+            diff = frac_diff @ cell
+
+        dist_sq = np.sum(diff * diff, axis=-1)
+        rows = np.arange(stop - start, dtype=np.int64)
+        dist_sq[rows, start + rows] = np.inf
+        min_dist_sq = min(min_dist_sq, float(dist_sq.min()))
+        if min_dist_sq == 0.0 or (
+            stop_dist_sq is not None and min_dist_sq < stop_dist_sq
+        ):
+            break
+
+    return float(np.sqrt(min_dist_sq))

deepmd/dpmodel/utils/lmdb_data.py

@@ -21,6 +21,9 @@
 import msgpack
 import numpy as np
 
+from deepmd.dpmodel.utils.dist_check import (
+    compute_min_pair_dist_single,
+)
 from deepmd.env import (
     GLOBAL_ENER_FLOAT_PRECISION,
     GLOBAL_NP_FLOAT_PRECISION,
@@ -597,6 +600,29 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             frame["natoms"] = fallback
             frame["real_natoms_vec"] = fallback
 
+        if "min_pair_dist" in self._data_requirements and "min_pair_dist" not in frame:
+            box = frame.get("box")
+            if box is not None and np.allclose(box, 0.0):
+                box = None
+            req = self._data_requirements["min_pair_dist"]
+            min_pair_dist = float(
+                req.get("default", 0.0)
+                if isinstance(req, dict)
+                else getattr(req, "default", 0.0)
+            )
+            frame["find_min_pair_dist"] = np.float32(1.0)
+            frame["min_pair_dist"] = np.array(
+                [
+                    compute_min_pair_dist_single(
+                        frame["coord"],
+                        box,
+                        frame["atype"],
+                        stop_below=min_pair_dist,
+                    )
+                ],
+                dtype=self._resolve_dtype("min_pair_dist"),
+            )
+
         # Add find_* flags for all data keys present in the frame.
         # Core structural keys and metadata are excluded — only label-like
         # and auxiliary data keys get find_* flags.

deepmd/dpmodel/utils/nlist.py

@@ -355,9 +355,16 @@ def extend_coord_with_ghosts(
         shift_idx = xp.take(xyz, xp.argsort(xp.linalg.vector_norm(xyz, axis=1)), axis=0)
         ns, _ = shift_idx.shape
         nall = ns * nloc
-        # shift_vec = xp.einsum("sd,fdk->fsk", shift_idx, cell)
-        shift_vec = xp.tensordot(shift_idx, cell, axes=([1], [1]))
-        shift_vec = xp.permute_dims(shift_vec, (1, 0, 2))
+        if array_api_compat.is_jax_namespace(xp):
+            # Avoid JAX internal errors in tensordot.
+            shift_vec = xp.sum(
+                shift_idx[xp.newaxis, :, :, xp.newaxis] * cell[:, xp.newaxis, :, :],
+                axis=2,
+            )
+        else:
+            # shift_vec = xp.einsum("sd,fdk->fsk", shift_idx, cell)
+            shift_vec = xp.tensordot(shift_idx, cell, axes=([1], [1]))
+            shift_vec = xp.permute_dims(shift_vec, (1, 0, 2))
         extend_coord = coord[:, None, :, :] + shift_vec[:, :, None, :]
         extend_atype = xp.tile(atype[:, :, xp.newaxis], (1, ns, 1))
         extend_aidx = xp.tile(aidx[:, :, xp.newaxis], (1, ns, 1))