style(dpmodel): enforce type annotations (deepmodeling#4953)

This PR fixes critical PyTorch JIT compilation errors that were
preventing model serialization and deployment. The main issue was that
PyTorch JIT cannot handle certain type annotations, specifically
`Union[np.ndarray, Any]` (the `Array` type) and `KeysView[str]` return
types.

## Problem

PyTorch JIT compilation was failing with errors like:
```
AssertionError: Unsupported annotation typing.Union[numpy.ndarray, typing.Any] could not be resolved because None could not be resolved.
```

This prevented users from using `torch.jit.script()` on trained models,
blocking deployment scenarios.

## Solution

- **Removed incompatible type annotations** from methods used by PyTorch
JIT compilation
- **Added `# noqa:ANNxxx` comments** to suppress ruff linting warnings
for missing type annotations
- **Fixed `KeysView[str]` return types** that caused JIT compilation
failures
- **Maintained type safety** elsewhere in the codebase where JIT
compatibility isn't required

## Files Changed

- `deepmd/dpmodel/utils/network.py`: Removed `Array` type annotations
from `call` methods and activation functions
- `deepmd/dpmodel/output_def.py`: Removed `KeysView[str]` return type
annotations from `keys` methods

## Validation

All PyTorch JIT tests now pass:
- ✅ TestEnergyModelSeA.test_jit
- ✅ TestDOSModelSeA.test_jit  
- ✅ TestEnergyModelDPA1.test_jit
- ✅ TestEnergyModelDPA2.test_jit
- ✅ TestEnergyModelHybrid.test_jit
- ✅ TestEnergyModelHybrid2.test_jit
- ✅ TestEnergyModelDPA2IntRcut.test_jit

Models can now be successfully JIT compiled for deployment while
maintaining full functionality and backward compatibility.

<!-- START COPILOT CODING AGENT TIPS -->
---

✨ Let Copilot coding agent [set things up for
you](https://github.com/deepmodeling/deepmd-kit/issues/new?title=✨+Set+up+Copilot+instructions&body=Configure%20instructions%20for%20this%20repository%20as%20documented%20in%20%5BBest%20practices%20for%20Copilot%20coding%20agent%20in%20your%20repository%5D%28https://gh.io/copilot-coding-agent-tips%29%2E%0A%0A%3COnboard%20this%20repo%3E&assignees=copilot)
— coding agent works faster and does higher quality work when set up for
your repo.

---------

Signed-off-by: Jinzhe Zeng <njzjz@qq.com>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 5 people

deepmd/dpmodel/array_api.py

@@ -1,14 +1,24 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Utilities for the array API."""
 
+from typing import (
+    Any,
+    Callable,
+    Optional,
+    Union,
+)
+
 import array_api_compat
 import numpy as np
 from packaging.version import (
     Version,
 )
 
+# Type alias for array_api compatible arrays
+Array = Union[np.ndarray, Any]  # Any to support JAX, PyTorch, etc. arrays
+
 
-def support_array_api(version: str) -> callable:
+def support_array_api(version: str) -> Callable:
     """Mark a function as supporting the specific version of the array API.
 
     Parameters
@@ -18,7 +28,7 @@ def support_array_api(version: str) -> callable:
 
     Returns
     -------
-    callable
+    Callable
         The decorated function
 
     Examples
@@ -28,7 +38,7 @@ def support_array_api(version: str) -> callable:
     ...     pass
     """
 
-    def set_version(func: callable) -> callable:
+    def set_version(func: Callable) -> Callable:
         func.array_api_version = version
         return func
 
@@ -39,15 +49,15 @@ def set_version(func: callable) -> callable:
 # but it hasn't been released yet
 # below is a pure Python implementation of take_along_axis
 # https://github.com/data-apis/array-api/issues/177#issuecomment-2093630595
-def xp_swapaxes(a, axis1, axis2):
+def xp_swapaxes(a: Array, axis1: int, axis2: int) -> Array:
     xp = array_api_compat.array_namespace(a)
     axes = list(range(a.ndim))
     axes[axis1], axes[axis2] = axes[axis2], axes[axis1]
     a = xp.permute_dims(a, axes)
     return a
 
 
-def xp_take_along_axis(arr, indices, axis):
+def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     xp = array_api_compat.array_namespace(arr)
     if Version(xp.__array_api_version__) >= Version("2024.12"):
         # see: https://github.com/data-apis/array-api-strict/blob/d086c619a58f35c38240592ef994aa19ca7beebc/array_api_strict/_indexing_functions.py#L30-L39
@@ -76,7 +86,7 @@ def xp_take_along_axis(arr, indices, axis):
     return xp_swapaxes(out, axis, -1)
 
 
-def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray:
+def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array:
     """Reduces all values from the src tensor to the indices specified in the index tensor."""
     # jax only
     if array_api_compat.is_jax_array(input):
@@ -94,7 +104,7 @@ def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray
         raise NotImplementedError("Only JAX arrays are supported.")
 
 
-def xp_add_at(x, indices, values):
+def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
     """Adds values to the specified indices of x in place or returns new x (for JAX)."""
     xp = array_api_compat.array_namespace(x, indices, values)
     if array_api_compat.is_numpy_array(x):
@@ -115,7 +125,7 @@ def xp_add_at(x, indices, values):
         return x
 
 
-def xp_bincount(x, weights=None, minlength=0):
+def xp_bincount(x: Array, weights: Optional[Array] = None, minlength: int = 0) -> Array:
     """Counts the number of occurrences of each value in x."""
     xp = array_api_compat.array_namespace(x)
     if array_api_compat.is_numpy_array(x) or array_api_compat.is_jax_array(x):

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,12 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import math
 from typing import (
+    Any,
     Optional,
 )
 
 import array_api_compat
 import numpy as np
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.common import (
     NativeOP,
     to_numpy_array,
@@ -42,7 +46,7 @@ def __init__(
         atom_exclude_types: list[int] = [],
         pair_exclude_types: list[tuple[int, int]] = [],
         rcond: Optional[float] = None,
-        preset_out_bias: Optional[dict[str, np.ndarray]] = None,
+        preset_out_bias: Optional[dict[str, Array]] = None,
     ) -> None:
         super().__init__()
         self.type_map = type_map
@@ -68,15 +72,15 @@ def init_out_stat(self) -> None:
         self.out_bias = out_bias_data
         self.out_std = out_std_data
 
-    def __setitem__(self, key, value) -> None:
+    def __setitem__(self, key: str, value: Array) -> None:
         if key in ["out_bias"]:
             self.out_bias = value
         elif key in ["out_std"]:
             self.out_std = value
         else:
             raise KeyError(key)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> Array:
         if key in ["out_bias"]:
             return self.out_bias
         elif key in ["out_std"]:
@@ -129,7 +133,7 @@ def atomic_output_def(self) -> FittingOutputDef:
         )
 
     def change_type_map(
-        self, type_map: list[str], model_with_new_type_stat=None
+        self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None
     ) -> None:
         """Change the type related params to new ones, according to `type_map` and the original one in the model.
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
@@ -147,13 +151,13 @@ def change_type_map(
 
     def forward_common_atomic(
         self,
-        extended_coord: np.ndarray,
-        extended_atype: np.ndarray,
-        nlist: np.ndarray,
-        mapping: Optional[np.ndarray] = None,
-        fparam: Optional[np.ndarray] = None,
-        aparam: Optional[np.ndarray] = None,
-    ) -> dict[str, np.ndarray]:
+        extended_coord: Array,
+        extended_atype: Array,
+        nlist: Array,
+        mapping: Optional[Array] = None,
+        fparam: Optional[Array] = None,
+        aparam: Optional[Array] = None,
+    ) -> dict[str, Array]:
         """Common interface for atomic inference.
 
         This method accept extended coordinates, extended atom typs, neighbor list,
@@ -223,13 +227,13 @@ def forward_common_atomic(
 
     def call(
         self,
-        extended_coord: np.ndarray,
-        extended_atype: np.ndarray,
-        nlist: np.ndarray,
-        mapping: Optional[np.ndarray] = None,
-        fparam: Optional[np.ndarray] = None,
-        aparam: Optional[np.ndarray] = None,
-    ) -> dict[str, np.ndarray]:
+        extended_coord: Array,
+        extended_atype: Array,
+        nlist: Array,
+        mapping: Optional[Array] = None,
+        fparam: Optional[Array] = None,
+        aparam: Optional[Array] = None,
+    ) -> dict[str, Array]:
         return self.forward_common_atomic(
             extended_coord,
             extended_atype,
@@ -264,9 +268,9 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel":
 
     def apply_out_stat(
         self,
-        ret: dict[str, np.ndarray],
-        atype: np.ndarray,
-    ):
+        ret: dict[str, Array],
+        atype: Array,
+    ) -> dict[str, Array]:
         """Apply the stat to each atomic output.
         The developer may override the method to define how the bias is applied
         to the atomic output of the model.
@@ -309,7 +313,7 @@ def _get_bias_index(
     def _fetch_out_stat(
         self,
         keys: list[str],
-    ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
+    ) -> tuple[dict[str, Array], dict[str, Array]]:
         ret_bias = {}
         ret_std = {}
         ntypes = self.get_ntypes()

deepmd/dpmodel/atomic_model/dipole_atomic_model.py

@@ -1,6 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import numpy as np
+from typing import (
+    Any,
+)
 
+from deepmd.dpmodel.array_api import (
+    Array,
+)
+from deepmd.dpmodel.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
+from deepmd.dpmodel.fitting.base_fitting import (
+    BaseFitting,
+)
 from deepmd.dpmodel.fitting.dipole_fitting import (
     DipoleFitting,
 )
@@ -11,7 +22,13 @@
 
 
 class DPDipoleAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs):
+    def __init__(
+        self,
+        descriptor: BaseDescriptor,
+        fitting: BaseFitting,
+        type_map: list[str],
+        **kwargs: Any,
+    ) -> None:
         if not isinstance(fitting, DipoleFitting):
             raise TypeError(
                 "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel"
@@ -20,8 +37,8 @@ def __init__(self, descriptor, fitting, type_map, **kwargs):
 
     def apply_out_stat(
         self,
-        ret: dict[str, np.ndarray],
-        atype: np.ndarray,
-    ):
+        ret: dict[str, Array],
+        atype: Array,
+    ) -> dict[str, Array]:
         # dipole not applying bias
         return ret

deepmd/dpmodel/atomic_model/dos_atomic_model.py

@@ -1,4 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+from deepmd.dpmodel.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
+from deepmd.dpmodel.fitting.base_fitting import (
+    BaseFitting,
+)
 from deepmd.dpmodel.fitting.dos_fitting import (
     DOSFittingNet,
 )
@@ -9,7 +19,13 @@
 
 
 class DPDOSAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs):
+    def __init__(
+        self,
+        descriptor: BaseDescriptor,
+        fitting: BaseFitting,
+        type_map: list[str],
+        **kwargs: Any,
+    ) -> None:
         if not isinstance(fitting, DOSFittingNet):
             raise TypeError(
                 "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel"

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
+    Any,
     Optional,
 )
 
-import numpy as np
-
+from deepmd.dpmodel.array_api import (
+    Array,
+)
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
@@ -41,10 +43,10 @@ class DPAtomicModel(BaseAtomicModel):
 
     def __init__(
         self,
-        descriptor,
-        fitting,
+        descriptor: BaseDescriptor,
+        fitting: BaseFitting,
         type_map: list[str],
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         super().__init__(type_map, **kwargs)
         self.type_map = type_map
@@ -65,7 +67,7 @@ def get_sel(self) -> list[int]:
         """Get the neighbor selection."""
         return self.descriptor.get_sel()
 
-    def set_case_embd(self, case_idx: int):
+    def set_case_embd(self, case_idx: int) -> None:
         """
         Set the case embedding of this atomic model by the given case_idx,
         typically concatenated with the output of the descriptor and fed into the fitting net.
@@ -125,13 +127,13 @@ def enable_compression(
 
     def forward_atomic(
         self,
-        extended_coord: np.ndarray,
-        extended_atype: np.ndarray,
-        nlist: np.ndarray,
-        mapping: Optional[np.ndarray] = None,
-        fparam: Optional[np.ndarray] = None,
-        aparam: Optional[np.ndarray] = None,
-    ) -> dict[str, np.ndarray]:
+        extended_coord: Array,
+        extended_atype: Array,
+        nlist: Array,
+        mapping: Optional[Array] = None,
+        fparam: Optional[Array] = None,
+        aparam: Optional[Array] = None,
+    ) -> dict[str, Array]:
         """Models' atomic predictions.
 
         Parameters
@@ -175,7 +177,7 @@ def forward_atomic(
         return ret
 
     def change_type_map(
-        self, type_map: list[str], model_with_new_type_stat=None
+        self, type_map: list[str], model_with_new_type_stat: Optional[Any] = None
     ) -> None:
         """Change the type related params to new ones, according to `type_map` and the original one in the model.
         If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
@@ -213,7 +215,7 @@ def serialize(self) -> dict:
     """The base fitting class."""
 
     @classmethod
-    def deserialize(cls, data) -> "DPAtomicModel":
+    def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel":
         data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
 from deepmd.dpmodel.fitting.ener_fitting import (
     EnergyFittingNet,
     InvarFitting,
@@ -10,7 +14,9 @@
 
 
 class DPEnergyAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs):
+    def __init__(
+        self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any
+    ) -> None:
         if not (
             isinstance(fitting, EnergyFittingNet) or isinstance(fitting, InvarFitting)
         ):