Merge branch 'devel' into fix-train-time

Author: caic99

.github/workflows/build_wheel.yml

@@ -44,7 +44,7 @@ jobs:
             platform_id: macosx_arm64
             dp_variant: cpu
           # win-64
-          - os: windows-2019
+          - os: windows-2022
             python: 311
             platform_id: win_amd64
             dp_variant: cpu
@@ -70,7 +70,7 @@ jobs:
           rm -rf .git
         if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.23
+        uses: pypa/cibuildwheel@v3.0
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.12
+    rev: v0.11.13
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -60,7 +60,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v20.1.5
+    rev: v20.1.6
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)

README.md

@@ -76,7 +76,7 @@ See [our v2 paper](https://doi.org/10.1063/5.0155600) for details of all feature
 #### v3
 
 - Multiple backends supported. Add PyTorch and JAX backends.
-- The DPA-2 and DPA-3 models.
+- The DPA2 and DPA3 models.
 - Plugin mechanisms for external models.
 
 See [our v3 paper](https://doi.org/10.1021/acs.jctc.5c00340) for details of all features until v3.0.

deepmd/dpmodel/array_api.py

@@ -60,7 +60,7 @@ def xp_take_along_axis(arr, indices, axis):
 
     shape = list(arr.shape)
     shape.pop(-1)
-    shape = [*shape, n]
+    shape = (*shape, n)
 
     arr = xp.reshape(arr, (-1,))
     if n != 0:

deepmd/dpmodel/descriptor/dpa1.py

@@ -520,7 +520,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim
@@ -1027,7 +1027,7 @@ def call(
                     xp.tile(
                         (xp.reshape(atype, (-1, 1)) * ntypes_with_padding), (1, nnei)
                     ),
-                    (-1),
+                    (-1,),
                 )
                 idx_j = xp.reshape(nei_type, (-1,))
                 # (nf x nl x nnei) x ng

deepmd/dpmodel/descriptor/dpa2.py

@@ -841,7 +841,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # repinit
         g1_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nframes, nall, self.tebd_dim),
         )
         g1_inp = g1_ext[:, :nloc, :]

deepmd/dpmodel/descriptor/dpa3.py

@@ -251,7 +251,7 @@ def deserialize(cls, data: dict) -> "RepFlowArgs":
 
 @BaseDescriptor.register("dpa3")
 class DescrptDPA3(NativeOP, BaseDescriptor):
-    r"""The DPA-3 descriptor[1]_.
+    r"""The DPA3 descriptor[1]_.
 
     Parameters
     ----------
@@ -562,12 +562,12 @@ def call(
         type_embedding = self.type_embedding.call()
         if self.use_loc_mapping:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0),
                 (nframes, nloc, self.tebd_dim),
             )
         else:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+                xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
                 (nframes, nall, self.tebd_dim),
             )
         node_ebd_inp = node_ebd_ext[:, :nloc, :]

deepmd/dpmodel/descriptor/repflows.py

@@ -526,7 +526,7 @@ def call(
         # edge_input, h2 = xp.split(dmatrix, [1], axis=-1)
         # nb x nloc x nnei x 1
         if self.edge_init_use_dist:
-            edge_input = xp.linalg.vector_norm(diff, axis=-1, keepdims=True)
+            edge_input = safe_for_vector_norm(diff, axis=-1, keepdims=True)
         else:
             edge_input = dmatrix[:, :, :, :1]
         h2 = dmatrix[:, :, :, 1:]
@@ -578,7 +578,8 @@ def call(
             # n_angle x 1
             a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask]
         else:
-            edge_index = angle_index = xp.zeros([1, 3], dtype=nlist.dtype)
+            edge_index = xp.zeros([2, 1], dtype=nlist.dtype)
+            angle_index = xp.zeros([3, 1], dtype=nlist.dtype)
 
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
@@ -622,7 +623,7 @@ def call(
                 edge_ebd,
                 h2,
                 sw,
-                owner=edge_index[:, 0],
+                owner=edge_index[0, :],
                 num_owner=nframes * nloc,
                 nb=nframes,
                 nloc=nloc,
@@ -1286,8 +1287,8 @@ def call(
         a_nlist: np.ndarray,  # nf x nloc x a_nnei
         a_nlist_mask: np.ndarray,  # nf x nloc x a_nnei
         a_sw: np.ndarray,  # switch func, nf x nloc x a_nnei
-        edge_index: np.ndarray,  # n_edge x 2
-        angle_index: np.ndarray,  # n_angle x 3
+        edge_index: np.ndarray,  # 2 x n_edge
+        angle_index: np.ndarray,  # 3 x n_angle
     ):
         """
         Parameters
@@ -1312,12 +1313,12 @@ def call(
             Masks of the neighbor list for angle. real nei 1 otherwise 0
         a_sw : nf x nloc x a_nnei
             Switch function for angle.
-        edge_index : Optional for dynamic sel, n_edge x 2
+        edge_index : Optional for dynamic sel, 2 x n_edge
             n2e_index : n_edge
                 Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i).
             n_ext2e_index : n_edge
                 Broadcast indices from extended node(j) to edge(ij).
-        angle_index : Optional for dynamic sel, n_angle x 3
+        angle_index : Optional for dynamic sel, 3 x n_angle
             n2a_index : n_angle
                 Broadcast indices from extended node(j) to angle(ijk).
             eij2a_index : n_angle
@@ -1362,11 +1363,11 @@ def call(
             assert (n_edge, 3) == h2.shape
         del a_nlist  # may be used in the future
 
-        n2e_index, n_ext2e_index = edge_index[:, 0], edge_index[:, 1]
+        n2e_index, n_ext2e_index = edge_index[0, :], edge_index[1, :]
         n2a_index, eij2a_index, eik2a_index = (
-            angle_index[:, 0],
-            angle_index[:, 1],
-            angle_index[:, 2],
+            angle_index[0, :],
+            angle_index[1, :],
+            angle_index[2, :],
         )
 
         # nb x nloc x nnei x n_dim [OR] n_edge x n_dim

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -358,7 +358,7 @@ def call(
         type_embedding = self.type_embedding.call()
         # nf x nall x tebd_dim
         atype_embd_ext = xp.reshape(
-            xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0),
+            xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim

deepmd/dpmodel/fitting/general_fitting.py

@@ -412,7 +412,7 @@ def _call_common(
                 )
             fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
             fparam = xp.tile(
-                xp.reshape(fparam, [nf, 1, self.numb_fparam]), (1, nloc, 1)
+                xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
             )
             xx = xp.concat(
                 [xx, fparam],
@@ -431,7 +431,7 @@ def _call_common(
                     f"get an input aparam of dim {aparam.shape[-1]}, "
                     f"which is not consistent with {self.numb_aparam}."
                 )
-            aparam = xp.reshape(aparam, [nf, nloc, self.numb_aparam])
+            aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
             aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]
             xx = xp.concat(
                 [xx, aparam],
@@ -446,7 +446,7 @@ def _call_common(
         if self.dim_case_embd > 0:
             assert self.case_embd is not None
             case_embd = xp.tile(
-                xp.reshape(self.case_embd[...], [1, 1, -1]), [nf, nloc, 1]
+                xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1)
             )
             xx = xp.concat(
                 [xx, case_embd],
@@ -465,7 +465,7 @@ def _call_common(
             )
             for type_i in range(self.ntypes):
                 mask = xp.tile(
-                    xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out)
+                    xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out)
                 )
                 atom_property = self.nets[(type_i,)](xx)
                 if self.remove_vaccum_contribution is not None and not (
@@ -485,10 +485,10 @@ def _call_common(
         outs += xp.reshape(
             xp.take(
                 xp.astype(self.bias_atom_e[...], outs.dtype),
-                xp.reshape(atype, [-1]),
+                xp.reshape(atype, (-1,)),
                 axis=0,
             ),
-            [nf, nloc, net_dim_out],
+            (nf, nloc, net_dim_out),
         )
         # nf x nloc
         exclude_mask = self.emask.build_type_exclude_mask(atype)