fix(paddle): Address additional code review feedback

- Changed zip strict=False to strict=True in nlist.py
- Changed ctx type from object to paddle.autograd.PyLayerContext in utils.py backward methods
- Changed TaskLoss.__init__ to accept **kwargs: Any and forward to return paddle.Tensor
- Added mae: bool | None = None parameter to TaskLoss.forward signature
- Changed all remaining object types to Any in frozen.py, model.py, ener.py, stat.py
- Fixed _can_be_converted_to_float return type from bool | None to bool
- Added Any imports where needed (ener.py, loss.py, frozen.py, model.py)

All code review feedback has been addressed.

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

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
@@ -54,7 +57,7 @@ def __init__(
         inference: bool = False,
         use_huber: bool = False,
         huber_delta: float = 0.01,
-        **kwargs: object,
+        **kwargs: Any,
     ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
 
@@ -547,7 +550,7 @@ def __init__(
         self,
         start_pref_h: float = 0.0,
         limit_pref_h: float = 0.0,
-        **kwargs: object,
+        **kwargs: Any,
     ) -> None:
         r"""Enable the layer to compute loss on hessian.
 

deepmd/pd/loss/loss.py

@@ -4,7 +4,7 @@
     abstractmethod,
 )
 from typing import (
-    NoReturn,
+    Any,
 )
 
 import paddle
@@ -18,7 +18,7 @@
 
 
 class TaskLoss(paddle.nn.Layer, ABC, make_plugin_registry("loss")):
-    def __init__(self, **kwargs) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         """Construct loss."""
         super().__init__()
 
@@ -29,7 +29,8 @@ def forward(
         label: dict[str, paddle.Tensor],
         natoms: int,
         learning_rate: float,
-    ) -> NoReturn:
+        mae: bool | None = None,
+    ) -> paddle.Tensor:
         """Return loss ."""
         raise NotImplementedError
 

deepmd/pd/model/model/__init__.py

@@ -65,7 +65,7 @@ def _get_standard_model_components(
     return descriptor, fitting, fitting_net["type"]
 
 
-def _can_be_converted_to_float(value: object) -> bool | None:
+def _can_be_converted_to_float(value: object) -> bool:
     try:
         float(value)
         return True

deepmd/pd/model/model/frozen.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import json
 from typing import (
+    Any,
     NoReturn,
 )
 
@@ -27,7 +28,7 @@ class FrozenModel(BaseModel):
         The path to the frozen model
     """
 
-    def __init__(self, model_file: str, **kwargs: object) -> None:
+    def __init__(self, model_file: str, **kwargs: Any) -> None:
         super().__init__(**kwargs)
         self.model_file = model_file
         if model_file.endswith(".json"):

deepmd/pd/model/model/model.py

@@ -4,6 +4,7 @@
     Callable,
 )
 from typing import (
+    Any,
     NoReturn,
 )
 
@@ -18,7 +19,7 @@
 
 
 class BaseModel(paddle.nn.Layer, make_base_model()):
-    def __init__(self, *args: object, **kwargs: object) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Construct a basic model for different tasks."""
         paddle.nn.Layer.__init__(self)
         self.model_def_script = ""

deepmd/pd/utils/nlist.py

@@ -404,7 +404,7 @@ def build_multiple_neighbor_list(
     rr.masked_fill(nlist_mask, float("inf"))
     nlist0 = nlist
     ret = {}
-    for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=False):
+    for rc, ns in zip(rcuts[::-1], nsels[::-1], strict=True):
         nlist0 = nlist0[:, :, :ns].masked_fill(rr[:, :, :ns] > rc, -1)
         ret[get_multiple_nlist_key(rc, ns)] = nlist0
     return ret

deepmd/pd/utils/stat.py

@@ -176,9 +176,7 @@ def _compute_model_predict(
         fparam = system.get("fparam", None)
         aparam = system.get("aparam", None)
 
-        def model_forward_auto_batch_size(
-            *args: object, **kwargs: object
-        ) -> paddle.Tensor:
+        def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> paddle.Tensor:
             return auto_batch_size.execute_all(
                 model_forward,
                 nframes,

deepmd/pd/utils/utils.py

@@ -148,7 +148,7 @@ def forward(
 
             @staticmethod
             def backward(
-                ctx: object, grad_grad_output: paddle.Tensor
+                ctx: paddle.autograd.PyLayerContext, grad_grad_output: paddle.Tensor
             ) -> tuple[paddle.Tensor, paddle.Tensor]:
                 (x, grad_output) = ctx.saved_tensor()
                 threshold = ctx.threshold