fix: type hint corrections only

This PR contains only type hint corrections extracted from PR njzjz#222:

- Return type corrections (NoReturn -> None, tuple[Array, Array] -> tuple[Array, Array, Array, Array, Array])
- Parameter type corrections (e_sel: int -> e_sel: int | list[int])
- Parameter name changes to match interface (extended_coord -> coord_ext)

Excludes:
- # type: ignore comments
- assert statements
- Variable renames (e_sel -> e_sel_list)
- Any logic changes

--
OpenClaw

Author: OpenClaw

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -152,7 +152,7 @@ def get_model_sels(self) -> list[int | list[int]]:
         """Get the sels for each individual models."""
         return [model.get_sel() for model in self.models]
 
-    def _sort_rcuts_sels(self) -> tuple[tuple[Array, Array], list[int]]:
+    def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]:
         # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut.
         zipped = sorted(
             zip(self.get_model_rcuts(), self.get_model_nsels(), strict=True),

deepmd/dpmodel/descriptor/descriptor.py

@@ -9,7 +9,6 @@
 )
 from typing import (
     Any,
-    NoReturn,
 )
 
 import array_api_compat
@@ -88,7 +87,7 @@ def compute_input_stats(
         self,
         merged: Callable[[], list[dict]] | list[dict],
         path: DPPath | None = None,
-    ) -> NoReturn:
+    ) -> None:
         """
         Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
 
@@ -113,7 +112,7 @@ def get_stats(self) -> dict[str, StatItem]:
 
     def share_params(
         self, base_class: Any, shared_level: Any, resume: bool = False
-    ) -> NoReturn:
+    ) -> None:
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
@@ -125,9 +124,9 @@ def share_params(
     def call(
         self,
         nlist: Array,
-        extended_coord: Array,
-        extended_atype: Array,
-        extended_atype_embd: Array | None = None,
+        coord_ext: Array,
+        atype_ext: Array,
+        atype_embd_ext: Array | None = None,
         mapping: Array | None = None,
         type_embedding: Array | None = None,
     ) -> Any:

deepmd/dpmodel/descriptor/dpa1.py

@@ -964,7 +964,7 @@ def call(
         atype_embd_ext: Array | None = None,
         mapping: Array | None = None,
         type_embedding: Array | None = None,
-    ) -> tuple[Array, Array]:
+    ) -> tuple[Array, Array, Array, Array, Array]:
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(

deepmd/dpmodel/descriptor/dpa2.py

@@ -808,7 +808,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
-    ) -> tuple[Array, Array]:
+    ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 
         Parameters

deepmd/dpmodel/descriptor/dpa3.py

@@ -531,7 +531,7 @@ def call(
         atype_ext: Array,
         nlist: Array,
         mapping: Array | None = None,
-    ) -> tuple[Array, Array]:
+    ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 
         Parameters

deepmd/dpmodel/descriptor/repflows.py

@@ -484,7 +484,8 @@ def call(
         atype_ext: Array,
         atype_embd_ext: Array | None = None,
         mapping: Array | None = None,
-    ) -> tuple[Array, Array]:
+        type_embedding: Array | None = None,
+    ) -> tuple[Array, Array, Array, Array, Array]:
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         nframes, nloc, nnei = nlist.shape
         nall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3
@@ -858,7 +859,7 @@ def __init__(
         self,
         e_rcut: float,
         e_rcut_smth: float,
-        e_sel: int,
+        e_sel: int | list[int],
         a_rcut: float,
         a_rcut_smth: float,
         a_sel: int,
@@ -1326,7 +1327,7 @@ def call(
         a_sw: Array,  # switch func, nf x nloc x a_nnei
         edge_index: Array,  # 2 x n_edge
         angle_index: Array,  # 3 x n_angle
-    ) -> tuple[Array, Array]:
+    ) -> tuple[Array, Array, Array]:
         """
         Parameters
         ----------