fix(paddle): Address code review feedback - Part 2

- Replaced all object types with Any in DPA1, DPA2, DPA3 descriptors
- Replaced object with Any in se_t_tebd.py, ener_model.py, make_model.py
- Replaced object with Any in task modules (ener.py, fitting.py, invar_fitting.py)
- Added Any import to all files that needed it
- Changed zip() strict=False to strict=True in se_a.py, transform_output.py, dataloader.py

All type annotations now use Any instead of object as requested in code review.

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

Author: Copilot

deepmd/pd/model/descriptor/dpa1.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -242,7 +245,7 @@ def __init__(
         use_tebd_bias: bool = False,
         type_map: list[str] | None = None,
         # not implemented
-        spin: object = None,
+        spin: Any = None,
         type: str | None = None,
     ) -> None:
         super().__init__()
@@ -398,7 +401,7 @@ def get_env_protection(self) -> float:
         return self.se_atten.get_env_protection()
 
     def share_params(
-        self, base_class: object, shared_level: int, resume: bool = False
+        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.
@@ -471,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: object = 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.
@@ -571,7 +574,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
             data["use_tebd_bias"] = True
         obj = cls(**data)
 
-        def t_cvt(xx: object) -> paddle.Tensor:
+        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(

deepmd/pd/model/descriptor/dpa2.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -149,7 +152,7 @@ def __init__(
         """
         super().__init__()
 
-        def init_subclass_params(sub_data: dict | object, sub_class: type) -> object:
+        def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
             if isinstance(sub_data, dict):
                 return sub_class(**sub_data)
             elif isinstance(sub_data, sub_class):
@@ -401,7 +404,7 @@ def get_env_protection(self) -> float:
         return self.repinit.get_env_protection()
 
     def share_params(
-        self, base_class: object, shared_level: int, resume: bool = False
+        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.
@@ -438,7 +441,7 @@ def share_params(
             raise NotImplementedError
 
     def change_type_map(
-        self, type_map: list[str], model_with_new_type_stat: object = 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.
@@ -674,7 +677,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         if obj.repinit.dim_out != obj.repformers.dim_in:
             obj.g1_shape_tranform = MLPLayer.deserialize(g1_shape_tranform)
 
-        def t_cvt(xx: object) -> paddle.Tensor:
+        def t_cvt(xx: Any) -> paddle.Tensor:
             return paddle.to_tensor(xx, dtype=obj.repinit.prec, place=env.DEVICE)
 
         # deserialize repinit

deepmd/pd/model/descriptor/dpa3.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -120,7 +123,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        def init_subclass_params(sub_data: dict | object, sub_class: type) -> object:
+        def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
             if isinstance(sub_data, dict):
                 return sub_class(**sub_data)
             elif isinstance(sub_data, sub_class):
@@ -303,7 +306,7 @@ def get_env_protection(self) -> float:
         return self.repflows.get_env_protection()
 
     def share_params(
-        self, base_class: object, shared_level: int, resume: bool = False
+        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.
@@ -332,7 +335,7 @@ def share_params(
             raise NotImplementedError
 
     def change_type_map(
-        self, type_map: list[str], model_with_new_type_stat: object = 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.
@@ -465,7 +468,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA3":
             type_embedding
         )
 
-        def t_cvt(xx: object) -> paddle.Tensor:
+        def t_cvt(xx: Any) -> paddle.Tensor:
             return paddle.to_tensor(xx, dtype=obj.repflows.prec, place=env.DEVICE)
 
         # deserialize repflow

deepmd/pd/model/descriptor/se_a.py

@@ -769,7 +769,7 @@ def forward(
                 self.filter_layers.networks,
                 self.compress_data,
                 self.compress_info,
-                strict=False,
+                strict=True,
             )
         ):
             if self.type_one_side:

deepmd/pd/model/descriptor/se_t_tebd.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -331,7 +334,7 @@ def get_stat_mean_and_stddev(self) -> tuple[paddle.Tensor, paddle.Tensor]:
         return self.se_ttebd.mean, self.se_ttebd.stddev
 
     def change_type_map(
-        self, type_map: list[str], model_with_new_type_stat: object | None = None
+        self, type_map: list[str], model_with_new_type_stat: Any | None = 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.

deepmd/pd/model/model/ener_model.py

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from typing import (
+    Any,
+)
+
 import paddle
 
 from deepmd.pd.model.atomic_model import (
@@ -25,8 +29,8 @@ class EnergyModel(DPModelCommon, DPEnergyModel_):
 
     def __init__(
         self,
-        *args: object,
-        **kwargs: object,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
         DPModelCommon.__init__(self)
         DPEnergyModel_.__init__(self, *args, **kwargs)

deepmd/pd/model/model/make_model.py

@@ -2,6 +2,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -65,10 +68,10 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type[BaseModel]:
     class CM(BaseModel):
         def __init__(
             self,
-            *args: object,
+            *args: Any,
             # underscore to prevent conflict with normal inputs
             atomic_model_: T_AtomicModel | None = None,
-            **kwargs: object,
+            **kwargs: Any,
         ) -> None:
             super().__init__(*args, **kwargs)
             if atomic_model_ is not None:

deepmd/pd/model/model/transform_output.py

@@ -122,7 +122,7 @@ def take_deriv(
     split_vv1 = paddle.split(vv1, [1] * size, axis=-1)
     split_svv1 = paddle.split(svv1, [1] * size, axis=-1)
     split_ff, split_avir = [], []
-    for vvi, svvi in zip(split_vv1, split_svv1, strict=False):
+    for vvi, svvi in zip(split_vv1, split_svv1, strict=True):
         # nf x nloc x 3, nf x nloc x 9
         ffi, aviri = task_deriv_one(
             vvi,

deepmd/pd/model/task/ener.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
 import logging
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -44,7 +47,7 @@ def __init__(
         mixed_types: bool = True,
         seed: int | list[int] | None = None,
         type_map: list[str] | None = None,
-        **kwargs: object,
+        **kwargs: Any,
     ) -> None:
         super().__init__(
             "energy",

deepmd/pd/model/task/fitting.py

@@ -6,6 +6,9 @@
 from collections.abc import (
     Callable,
 )
+from typing import (
+    Any,
+)
 
 import numpy as np
 import paddle
@@ -51,7 +54,7 @@
 class Fitting(paddle.nn.Layer, BaseFitting):
     # plugin moved to BaseFitting
 
-    def __new__(cls, *args: object, **kwargs: object) -> "Fitting":
+    def __new__(cls, *args: Any, **kwargs: Any) -> "Fitting":
         if cls is Fitting:
             return BaseFitting.__new__(BaseFitting, *args, **kwargs)
         return super().__new__(cls)
@@ -244,7 +247,7 @@ def __init__(
         type_map: list[str] | None = None,
         use_aparam_as_mask: bool = False,
         default_fparam: list[float] | None = None,
-        **kwargs: object,
+        **kwargs: Any,
     ) -> None:
         super().__init__()
         self.var_name = var_name