refact(pt_expt): add decorator to simplify the module (#5213)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Introduced a public adapter to expose DP classes as PyTorch modules
via a decorator-based API.

* **Refactor**
* Converted multiple descriptor, network, and utility classes to
decorator-driven PyTorch integration, simplifying initialization and
attribute handling.

* **Breaking Changes**
* Several descriptor forward signatures expanded to accept extended
topology/embedding inputs and now return expanded tuples (update call
sites accordingly).
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <njzjz@qq.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 3 people

deepmd/pt_expt/common.py

@@ -17,6 +17,9 @@
 from collections.abc import (
     Callable,
 )
+from functools import (
+    wraps,
+)
 from typing import (
     Any,
     overload,
@@ -292,6 +295,46 @@ def to_torch_array(array: Any) -> torch.Tensor | None:
     return torch.as_tensor(array, device=env.DEVICE)
 
 
+def torch_module(
+    module: type[NativeOP],
+) -> type[torch.nn.Module]:
+    """Convert a NativeOP to a torch.nn.Module.
+
+    Parameters
+    ----------
+    module : type[NativeOP]
+        The NativeOP to convert.
+
+    Returns
+    -------
+    type[torch.nn.Module]
+        The torch.nn.Module.
+
+    Examples
+    --------
+    >>> @torch_module
+    ... class MyModule(NativeOP):
+    ...     pass
+    """
+
+    @wraps(module, updated=())
+    class TorchModule(module, torch.nn.Module):
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            torch.nn.Module.__init__(self)
+            module.__init__(self, *args, **kwargs)
+
+        def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
+            return torch.nn.Module.__call__(self, *args, **kwargs)
+
+        def __setattr__(self, name: str, value: Any) -> None:
+            handled, value = dpmodel_setattr(self, name, value)
+            if not handled:
+                super().__setattr__(name, value)
+
+    return TorchModule
+
+
 # Import utils to trigger dpmodel→pt_expt converter registrations
 # This must happen after the functions above are defined to avoid circular imports
 def _ensure_registrations() -> None:

deepmd/pt_expt/descriptor/se_e2_a.py

@@ -1,13 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
 import torch
 
 from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeAArrayAPI as DescrptSeADP
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
+    torch_module,
 )
 from deepmd.pt_expt.descriptor.base_descriptor import (
     BaseDescriptor,
@@ -16,20 +13,8 @@
 
 @BaseDescriptor.register("se_e2_a_expt")
 @BaseDescriptor.register("se_a_expt")
-class DescrptSeA(DescrptSeADP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        DescrptSeADP.__init__(self, *args, **kwargs)
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
-        return torch.nn.Module.__call__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
-
+@torch_module
+class DescrptSeA(DescrptSeADP):
     def forward(
         self,
         extended_coord: torch.Tensor,

deepmd/pt_expt/descriptor/se_r.py

@@ -1,13 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
 import torch
 
 from deepmd.dpmodel.descriptor.se_r import DescrptSeR as DescrptSeRDP
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
+    torch_module,
 )
 from deepmd.pt_expt.descriptor.base_descriptor import (
     BaseDescriptor,
@@ -16,20 +13,8 @@
 
 @BaseDescriptor.register("se_e2_r_expt")
 @BaseDescriptor.register("se_r_expt")
-class DescrptSeR(DescrptSeRDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        DescrptSeRDP.__init__(self, *args, **kwargs)
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
-        return torch.nn.Module.__call__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
-
+@torch_module
+class DescrptSeR(DescrptSeRDP):
     def forward(
         self,
         extended_coord: torch.Tensor,

deepmd/pt_expt/descriptor/se_t.py

@@ -1,13 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
 import torch
 
 from deepmd.dpmodel.descriptor.se_t import DescrptSeT as DescrptSeTDP
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
+    torch_module,
 )
 from deepmd.pt_expt.descriptor.base_descriptor import (
     BaseDescriptor,
@@ -17,20 +14,8 @@
 @BaseDescriptor.register("se_e3_expt")
 @BaseDescriptor.register("se_at_expt")
 @BaseDescriptor.register("se_a_3be_expt")
-class DescrptSeT(DescrptSeTDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        DescrptSeTDP.__init__(self, *args, **kwargs)
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
-        return torch.nn.Module.__call__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
-
+@torch_module
+class DescrptSeT(DescrptSeTDP):
     def forward(
         self,
         extended_coord: torch.Tensor,

deepmd/pt_expt/descriptor/se_t_tebd.py

@@ -1,34 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
 import torch
 
 from deepmd.dpmodel.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdDP
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
+    torch_module,
 )
 from deepmd.pt_expt.descriptor.base_descriptor import (
     BaseDescriptor,
 )
 
 
 @BaseDescriptor.register("se_e3_tebd_expt")
-class DescrptSeTTebd(DescrptSeTTebdDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        DescrptSeTTebdDP.__init__(self, *args, **kwargs)
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        # Ensure torch.nn.Module.__call__ drives forward() for export/tracing.
-        return torch.nn.Module.__call__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
-
+@torch_module
+class DescrptSeTTebd(DescrptSeTTebdDP):
     def forward(
         self,
         extended_coord: torch.Tensor,

deepmd/pt_expt/descriptor/se_t_tebd_block.py

@@ -1,28 +1,41 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
 import torch
 
 from deepmd.dpmodel.descriptor.se_t_tebd import (
     DescrptBlockSeTTebd as DescrptBlockSeTTebdDP,
 )
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
     register_dpmodel_mapping,
+    torch_module,
 )
 
 
-class DescrptBlockSeTTebd(DescrptBlockSeTTebdDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        DescrptBlockSeTTebdDP.__init__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
+@torch_module
+class DescrptBlockSeTTebd(DescrptBlockSeTTebdDP):
+    def forward(
+        self,
+        nlist: torch.Tensor,
+        coord_ext: torch.Tensor,
+        atype_ext: torch.Tensor,
+        atype_embd_ext: torch.Tensor | None = None,
+        mapping: torch.Tensor | None = None,
+        type_embedding: torch.Tensor | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        return self.call(
+            nlist,
+            coord_ext,
+            atype_ext,
+            atype_embd_ext=atype_embd_ext,
+            mapping=mapping,
+            type_embedding=type_embedding,
+        )
 
 
 register_dpmodel_mapping(

deepmd/pt_expt/utils/exclude_mask.py

@@ -1,27 +1,17 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
 
-import torch
 
 from deepmd.dpmodel.utils.exclude_mask import AtomExcludeMask as AtomExcludeMaskDP
 from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask as PairExcludeMaskDP
 from deepmd.pt_expt.common import (
-    dpmodel_setattr,
     register_dpmodel_mapping,
+    torch_module,
 )
 
 
-class AtomExcludeMask(AtomExcludeMaskDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        AtomExcludeMaskDP.__init__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
+@torch_module
+class AtomExcludeMask(AtomExcludeMaskDP):
+    pass
 
 
 register_dpmodel_mapping(
@@ -30,15 +20,9 @@ def __setattr__(self, name: str, value: Any) -> None:
 )
 
 
-class PairExcludeMask(PairExcludeMaskDP, torch.nn.Module):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
-        PairExcludeMaskDP.__init__(self, *args, **kwargs)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        handled, value = dpmodel_setattr(self, name, value)
-        if not handled:
-            super().__setattr__(name, value)
+@torch_module
+class PairExcludeMask(PairExcludeMaskDP):
+    pass
 
 
 register_dpmodel_mapping(

deepmd/pt_expt/utils/network.py

@@ -21,6 +21,7 @@
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
     to_torch_array,
+    torch_module,
 )
 
 
@@ -37,6 +38,7 @@ def __array__(self, dtype: Any | None = None) -> np.ndarray:
         return arr.astype(dtype)
 
 
+# do not apply torch_module until its setattr working to register parameters
 class NativeLayer(NativeLayerDP, torch.nn.Module):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         torch.nn.Module.__init__(self)
@@ -78,15 +80,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.call(x)
 
 
-class NativeNet(make_multilayer_network(NativeLayer, NativeOP), torch.nn.Module):
+@torch_module
+class NativeNet(make_multilayer_network(NativeLayer, NativeOP)):
     def __init__(self, layers: list[dict] | None = None) -> None:
-        torch.nn.Module.__init__(self)
         super().__init__(layers)
         self.layers = torch.nn.ModuleList(self.layers)
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        return torch.nn.Module.__call__(self, *args, **kwargs)
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.call(x)
 
@@ -136,15 +135,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 )
 
 
-class NetworkCollection(NetworkCollectionDP, torch.nn.Module):
+@torch_module
+class NetworkCollection(NetworkCollectionDP):
     NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = {
         "network": NativeNet,
         "embedding_network": EmbeddingNet,
         "fitting_network": FittingNet,
     }
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        torch.nn.Module.__init__(self)
         self._module_networks = torch.nn.ModuleDict()
         super().__init__(*args, **kwargs)