style(paddle): Add comprehensive type hints to Paddle backend and enable ANN rules

deepmd/pd/entrypoints/main.py

@@ -574,7 +574,7 @@ def change_bias(
     log.info(f"Saved model to {output_path}")
 
 
-def main(args: list[str] | argparse.Namespace | None = None):
+def main(args: list[str] | argparse.Namespace | None = None) -> None:
     if not isinstance(args, argparse.Namespace):
         FLAGS = parse_args(args=args)
     else:

deepmd/pd/infer/deep_eval.py

@@ -474,7 +474,7 @@ def _eval_model(
         fparam: np.ndarray | None,
         aparam: np.ndarray | None,
         request_defs: list[OutputVariableDef],
-    ):
+    ) -> tuple[np.ndarray, ...]:
         if not self.static_model:
             model = self.dp.to(DEVICE)
         prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PD_FLOAT_PRECISION]]

deepmd/pd/infer/inference.py

@@ -23,9 +23,9 @@
 class Tester:
     def __init__(
         self,
-        model_ckpt,
-        head=None,
-    ):
+        model_ckpt: str,
+        head: str | None = None,
+    ) -> None:
         """Construct a DeePMD tester.
 
         Args:

deepmd/pd/loss/ener.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
 
 import paddle
 import paddle.nn.functional as F
@@ -20,7 +23,9 @@
 )
 
 
-def custom_huber_loss(predictions, targets, delta=1.0):
+def custom_huber_loss(
+    predictions: paddle.Tensor, targets: paddle.Tensor, delta: float = 1.0
+) -> paddle.Tensor:
     error = targets - predictions
     abs_error = paddle.abs(error)
     quadratic_loss = 0.5 * paddle.pow(error, 2)
@@ -32,13 +37,13 @@ def custom_huber_loss(predictions, targets, delta=1.0):
 class EnergyStdLoss(TaskLoss):
     def __init__(
         self,
-        starter_learning_rate=1.0,
-        start_pref_e=0.0,
-        limit_pref_e=0.0,
-        start_pref_f=0.0,
-        limit_pref_f=0.0,
-        start_pref_v=0.0,
-        limit_pref_v=0.0,
+        starter_learning_rate: float = 1.0,
+        start_pref_e: float = 0.0,
+        limit_pref_e: float = 0.0,
+        start_pref_f: float = 0.0,
+        limit_pref_f: float = 0.0,
+        start_pref_v: float = 0.0,
+        limit_pref_v: float = 0.0,
         start_pref_ae: float = 0.0,
         limit_pref_ae: float = 0.0,
         start_pref_pf: float = 0.0,
@@ -49,10 +54,10 @@ def __init__(
         limit_pref_gf: float = 0.0,
         numb_generalized_coord: int = 0,
         use_l1_all: bool = False,
-        inference=False,
-        use_huber=False,
-        huber_delta=0.01,
-        **kwargs,
+        inference: bool = False,
+        use_huber: bool = False,
+        huber_delta: float = 0.01,
+        **kwargs: Any,
     ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
 
@@ -146,7 +151,15 @@ def __init__(
                 "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
             )
 
-    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+    def forward(
+        self,
+        input_dict: dict[str, paddle.Tensor],
+        model: paddle.nn.Layer,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool = False,
+    ) -> tuple[dict[str, paddle.Tensor], paddle.Tensor, dict[str, paddle.Tensor]]:
         """Return loss on energy and force.
 
         Parameters
@@ -535,10 +548,10 @@ def deserialize(cls, data: dict) -> "TaskLoss":
 class EnergyHessianStdLoss(EnergyStdLoss):
     def __init__(
         self,
-        start_pref_h=0.0,
-        limit_pref_h=0.0,
-        **kwargs,
-    ):
+        start_pref_h: float = 0.0,
+        limit_pref_h: float = 0.0,
+        **kwargs: Any,
+    ) -> None:
         r"""Enable the layer to compute loss on hessian.
 
         Parameters
@@ -556,7 +569,15 @@ def __init__(
         self.start_pref_h = start_pref_h
         self.limit_pref_h = limit_pref_h
 
-    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+    def forward(
+        self,
+        input_dict: dict[str, paddle.Tensor],
+        model: paddle.nn.Module,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool = False,
+    ) -> tuple[dict[str, paddle.Tensor], paddle.Tensor, dict[str, paddle.Tensor]]:
         model_pred, loss, more_loss = super().forward(
             input_dict, model, label, natoms, learning_rate, mae=mae
         )

deepmd/pd/loss/loss.py

@@ -3,6 +3,9 @@
     ABC,
     abstractmethod,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -15,11 +18,19 @@
 
 
 class TaskLoss(paddle.nn.Layer, ABC, make_plugin_registry("loss")):
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any) -> None:
         """Construct loss."""
         super().__init__()
 
-    def forward(self, input_dict, model, label, natoms, learning_rate):
+    def forward(
+        self,
+        input_dict: dict[str, paddle.Tensor],
+        model: paddle.nn.Module,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool | None = None,
+    ) -> paddle.Tensor:
         """Return loss ."""
         raise NotImplementedError
 

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -578,7 +578,7 @@ def _default_bias(self) -> paddle.Tensor:
             device=device
         )
 
-    def _default_std(self):
+    def _default_std(self) -> paddle.Tensor:
         ntypes = self.get_ntypes()
         return paddle.ones([self.n_out, ntypes, self.max_out_size], dtype=dtype).to(
             device=device
@@ -626,7 +626,7 @@ def _store_out_stat(
         paddle.assign(out_bias_data, self.out_bias)
         paddle.assign(out_std_data, self.out_std)
 
-    def get_ntypes(self):
+    def get_ntypes(self) -> int:
         return len(self.type_map)
 
     def get_buffer_ntypes(self) -> paddle.Tensor:

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -385,7 +385,7 @@ def compute_or_load_stat(
             stat_file_path /= " ".join(self.type_map)
 
         @functools.lru_cache
-        def wrapped_sampler():
+        def wrapped_sampler() -> list[dict]:
             sampled = sampled_func()
             if self.pair_excl is not None:
                 pair_exclude_types = self.pair_excl.get_exclude_types()

deepmd/pd/model/atomic_model/energy_atomic_model.py

@@ -1,4 +1,14 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+from deepmd.pd.model.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
+from deepmd.pd.model.task.base_fitting import (
+    BaseFitting,
+)
 from deepmd.pd.model.task.ener import (
     EnergyFittingNet,
     InvarFitting,
@@ -10,7 +20,13 @@
 
 
 class DPEnergyAtomicModel(DPAtomicModel):
-    def __init__(self, descriptor, fitting, type_map, **kwargs):
+    def __init__(
+        self,
+        descriptor: BaseDescriptor,
+        fitting: BaseFitting,
+        type_map: list[str],
+        **kwargs: Any,
+    ) -> None:
         assert isinstance(fitting, EnergyFittingNet) or isinstance(
             fitting, InvarFitting
         )

deepmd/pd/model/descriptor/descriptor.py

@@ -8,6 +8,7 @@
     Callable,
 )
 from typing import (
+    Any,
     NoReturn,
 )
 
@@ -43,7 +44,7 @@ class DescriptorBlock(paddle.nn.Layer, ABC, make_plugin_registry("DescriptorBloc
 
     local_cluster = False
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args: Any, **kwargs: Any) -> "DescriptorBlock":
         if cls is DescriptorBlock:
             try:
                 descrpt_type = kwargs["type"]
@@ -126,7 +127,9 @@ def get_stats(self) -> dict[str, StatItem]:
         """Get the statistics of the descriptor."""
         raise NotImplementedError
 
-    def share_params(self, base_class, shared_level, resume=False) -> None:
+    def share_params(
+        self, base_class: Any, shared_level: int, resume: bool = False
+    ) -> None:
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
@@ -180,7 +183,7 @@ def forward(
         extended_atype_embd: paddle.Tensor | None = None,
         mapping: paddle.Tensor | None = None,
         type_embedding: paddle.Tensor | None = None,
-    ):
+    ) -> paddle.Tensor:
         """Calculate DescriptorBlock."""
         pass
 
@@ -194,14 +197,16 @@ def need_sorted_nlist_for_lower(self) -> bool:
 
 
 def make_default_type_embedding(
-    ntypes,
-):
+    ntypes: int,
+) -> tuple[TypeEmbedNet, dict]:
     aux = {}
     aux["tebd_dim"] = 8
     return TypeEmbedNet(ntypes, aux["tebd_dim"]), aux
 
 
-def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None:
+def extend_descrpt_stat(
+    des: Any, type_map: list[str], des_with_stat: Any = None
+) -> None:
     r"""
     Extend the statistics of a descriptor block with types from newly provided `type_map`.
 

deepmd/pd/model/descriptor/dpa1.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -228,8 +231,8 @@ def __init__(
         exclude_types: list[tuple[int, int]] = [],
         env_protection: float = 0.0,
         scaling_factor: int = 1.0,
-        normalize=True,
-        temperature=None,
+        normalize: bool = True,
+        temperature: float | None = None,
         concat_output_tebd: bool = True,
         trainable: bool = True,
         trainable_ln: bool = True,
@@ -242,7 +245,7 @@ def __init__(
         use_tebd_bias: bool = False,
         type_map: list[str] | None = None,
         # not implemented
-        spin=None,
+        spin: Any = None,
         type: str | None = None,
     ) -> None:
         super().__init__()
@@ -397,7 +400,9 @@ def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
 
-    def share_params(self, base_class, shared_level, resume=False) -> None:
+    def share_params(
+        self, base_class: Any, shared_level: int, resume: bool = False
+    ) -> None:
         """
         Share the parameters of self to the base_class with shared_level during multitask training.
         If not start from checkpoint (resume is False),
@@ -425,18 +430,18 @@ def share_params(self, base_class, shared_level, resume=False) -> None:
             raise NotImplementedError
 
     @property
-    def dim_out(self):
+    def dim_out(self) -> int:
         return self.get_dim_out()
 
     @property
-    def dim_emb(self):
+    def dim_emb(self) -> int:
         return self.get_dim_emb()
 
     def compute_input_stats(
         self,
         merged: Callable[[], list[dict]] | list[dict],
         path: DPPath | None = None,
-    ):
+    ) -> None:
         """
         Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
 
@@ -469,7 +474,7 @@ def get_stat_mean_and_stddev(self) -> tuple[paddle.Tensor, paddle.Tensor]:
         return self.se_atten.mean, self.se_atten.stddev
 
     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: 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.
@@ -569,7 +574,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
             data["use_tebd_bias"] = True
         obj = cls(**data)
 
-        def t_cvt(xx):
+        def t_cvt(xx: Any) -> paddle.Tensor:
             return paddle.to_tensor(xx, dtype=obj.se_atten.prec).to(device=env.DEVICE)
 
         obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize(
@@ -620,7 +625,7 @@ def forward(
         nlist: paddle.Tensor,
         mapping: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
-    ):
+    ) -> paddle.Tensor:
         """Compute the descriptor.
 
         Parameters