feat(paddle): Add type hints to Paddle backend and enable ANN rule for entry points

Co-authored-by: njzjz <[email protected]>

Author: Copilot

deepmd/pd/entrypoints/main.py

@@ -7,6 +7,7 @@
     Path,
 )
 from typing import (
+    Any,
     Optional,
     Union,
 )
@@ -80,15 +81,15 @@
 
 
 def get_trainer(
-    config,
-    init_model=None,
-    restart_model=None,
-    finetune_model=None,
-    force_load=False,
-    init_frz_model=None,
-    shared_links=None,
-    finetune_links=None,
-):
+    config: dict[str, Any],
+    init_model: Optional[str] = None,
+    restart_model: Optional[str] = None,
+    finetune_model: Optional[str] = None,
+    force_load: bool = False,
+    init_frz_model: Optional[str] = None,
+    shared_links: Optional[dict[str, Any]] = None,
+    finetune_links: Optional[dict[str, Any]] = None,
+) -> training.Trainer:
     multi_task = "model_dict" in config.get("model", {})
 
     # Initialize DDP
@@ -98,8 +99,11 @@ def get_trainer(
         fleet.init(is_collective=True)
 
     def prepare_trainer_input_single(
-        model_params_single, data_dict_single, rank=0, seed=None
-    ):
+        model_params_single: dict[str, Any],
+        data_dict_single: dict[str, Any],
+        rank: int = 0,
+        seed: Optional[int] = None,
+    ) -> tuple[Any, Any, Any, Optional[Any]]:
         training_dataset_params = data_dict_single["training_data"]
         validation_dataset_params = data_dict_single.get("validation_data", None)
         validation_systems = (
@@ -535,7 +539,7 @@ def change_bias(
     log.info(f"Saved model to {output_path}")
 
 
-def main(args: Optional[Union[list[str], argparse.Namespace]] = None):
+def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
     if not isinstance(args, argparse.Namespace):
         FLAGS = parse_args(args=args)
     else:

deepmd/pd/train/training.py

@@ -11,6 +11,8 @@
 )
 from typing import (
     Any,
+    Optional,
+    Union,
 )
 
 import numpy as np
@@ -86,16 +88,16 @@ class Trainer:
     def __init__(
         self,
         config: dict[str, Any],
-        training_data,
-        stat_file_path=None,
-        validation_data=None,
-        init_model=None,
-        restart_model=None,
-        finetune_model=None,
-        force_load=False,
-        shared_links=None,
-        finetune_links=None,
-        init_frz_model=None,
+        training_data: Any,
+        stat_file_path: Optional[Union[str, Path]] = None,
+        validation_data: Optional[Any] = None,
+        init_model: Optional[str] = None,
+        restart_model: Optional[str] = None,
+        finetune_model: Optional[str] = None,
+        force_load: bool = False,
+        shared_links: Optional[dict[str, Any]] = None,
+        finetune_links: Optional[dict[str, Any]] = None,
+        init_frz_model: Optional[str] = None,
     ) -> None:
         """Construct a DeePMD trainer.
 
@@ -1057,7 +1059,7 @@ def log_loss_valid(_task_key="Default"):
                 "files, which can be viewd in NVIDIA Nsight Systems software"
             )
 
-    def save_model(self, save_path, lr=0.0, step=0) -> None:
+    def save_model(self, save_path: str, lr: float = 0.0, step: int = 0) -> None:
         module = (
             self.wrapper._layers
             if dist.is_available() and dist.is_initialized()
@@ -1079,7 +1081,9 @@ def save_model(self, save_path, lr=0.0, step=0) -> None:
             checkpoint_files.sort(key=lambda x: x.stat().st_mtime)
             checkpoint_files[0].unlink()
 
-    def get_data(self, is_train=True, task_key="Default"):
+    def get_data(
+        self, is_train: bool = True, task_key: str = "Default"
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
         if not self.multi_task:
             if is_train:
                 try:
@@ -1155,7 +1159,9 @@ def get_data(self, is_train=True, task_key="Default"):
         log_dict["sid"] = batch_data["sid"]
         return input_dict, label_dict, log_dict
 
-    def print_header(self, fout, train_results, valid_results) -> None:
+    def print_header(
+        self, fout: Any, train_results: dict[str, Any], valid_results: dict[str, Any]
+    ) -> None:
         train_keys = sorted(train_results.keys())
         print_str = ""
         print_str += "# {:5s}".format("step")
@@ -1187,7 +1193,12 @@ def print_header(self, fout, train_results, valid_results) -> None:
         fout.flush()
 
     def print_on_training(
-        self, fout, step_id, cur_lr, train_results, valid_results
+        self,
+        fout: Any,
+        step_id: int,
+        cur_lr: float,
+        train_results: dict[str, Any],
+        valid_results: dict[str, Any],
     ) -> None:
         train_keys = sorted(train_results.keys())
         print_str = ""

deepmd/pd/utils/dataloader.py

@@ -12,6 +12,10 @@
 from threading import (
     Thread,
 )
+from typing import (
+    Optional,
+    Union,
+)
 
 import h5py
 import numpy as np
@@ -53,7 +57,7 @@
 # paddle.multiprocessing.set_sharing_strategy("file_system")
 
 
-def setup_seed(seed):
+def setup_seed(seed: Union[int, list, tuple]) -> None:
     if isinstance(seed, (list, tuple)):
         mixed_seed = mix_entropy(seed)
     else:
@@ -82,12 +86,12 @@ class DpLoaderSet(Dataset):
 
     def __init__(
         self,
-        systems,
-        batch_size,
-        type_map,
-        seed=None,
-        shuffle=True,
-    ):
+        systems: Union[str, list[str]],
+        batch_size: int,
+        type_map: list[str],
+        seed: Optional[int] = None,
+        shuffle: bool = True,
+    ) -> None:
         if seed is not None:
             setup_seed(seed)
         if isinstance(systems, str):

deepmd/pd/utils/utils.py

@@ -83,7 +83,7 @@ def silut_double_backward(
 
 
 class SiLUTScript(paddle.nn.Layer):
-    def __init__(self, threshold: float = 3.0):
+    def __init__(self, threshold: float = 3.0) -> None:
         super().__init__()
         self.threshold = threshold
 
@@ -95,7 +95,7 @@ def __init__(self, threshold: float = 3.0):
         self.const_val = float(threshold * sigmoid_threshold)
         self.get_script_code()
 
-    def get_script_code(self):
+    def get_script_code(self) -> None:
         silut_forward_script = paddle.jit.to_static(silut_forward, full_graph=True)
         silut_backward_script = paddle.jit.to_static(silut_backward, full_graph=True)
         silut_double_backward_script = paddle.jit.to_static(
@@ -142,12 +142,12 @@ def backward(ctx, grad_grad_output):
 
         self.SiLUTFunction = SiLUTFunction
 
-    def forward(self, x):
+    def forward(self, x: paddle.Tensor) -> paddle.Tensor:
         return self.SiLUTFunction.apply(x, self.threshold, self.slope, self.const_val)
 
 
 class SiLUT(paddle.nn.Layer):
-    def __init__(self, threshold=3.0):
+    def __init__(self, threshold: float = 3.0) -> None:
         super().__init__()
 
         def sigmoid(x):

pyproject.toml

@@ -426,7 +426,16 @@ runtime-evaluated-base-classes = ["torch.nn.Module"]
 "deepmd/tf/**" = ["TID253", "ANN"]
 "deepmd/pt/**" = ["TID253", "ANN"]
 "deepmd/jax/**" = ["TID253", "ANN"]
-"deepmd/pd/**" = ["TID253", "ANN"]
+# Paddle backend: Gradually enabling ANN rule
+# Completed files with full type annotations:
+"deepmd/pd/entrypoints/main.py" = ["TID253"]  # ✅ Fully typed
+# TODO: Complete type hints and remove ANN exclusion for remaining files:
+"deepmd/pd/train/**" = ["TID253", "ANN"]     # 🚧 Partial progress
+"deepmd/pd/utils/**" = ["TID253", "ANN"]     # 🚧 Partial progress
+"deepmd/pd/loss/**" = ["TID253", "ANN"]      # ❌ Not started
+"deepmd/pd/model/**" = ["TID253", "ANN"]     # ❌ Not started
+"deepmd/pd/infer/**" = ["TID253", "ANN"]     # ❌ Not started
+"deepmd/pd/cxx_op.py" = ["ANN"]              # ❌ Not started
 "deepmd/dpmodel/**" = ["ANN"]
 "source/**" = ["ANN"]
 "source/tests/tf/**" = ["TID253", "ANN"]