feat(pt_expt): auto-generate forward/forward_lower in torch_module decorator (#5246)

Copilot · njzjz · web-flow · commit 4ddc37d8c829 · 2026-02-18T08:36:24.000Z
- [x] Understand the current implementation of `torch_module` decorator - [x] Identify all classes using `@torch_module` with manual `forward()` methods - [x] Modify `torch_module` decorator to auto-generate `forward()` from `call()` - [x] Modify `torch_module` decorator to auto-generate `forward_lower()` from `call_lower()` - [x] Remove boilerplate `forward()` methods from descriptor classes (DescrptSeA, DescrptSeR, DescrptSeT, DescrptSeTTebd, DescrptBlockSeTTebd) - [x] Remove boilerplate `forward()` methods from fitting classes (InvarFitting, EnergyFittingNet) - [x] Remove boilerplate `forward()` method from NativeNet - [x] Run linting and formatting - [x] Create tests to verify auto-generated methods work correctly - [x] Run targeted tests to verify changes (63 pt_expt tests pass) - [x] Manually verify forward methods work with real data - [x] Request code review and address feedback - [x] Update docstring to document auto-generation behavior - [x] Update tests to use module(...) invocation path - [x] Ready for final approval  <details> <summary>Original prompt</summary> > > ---- > > *This section details on the original issue you should resolve* > > <issue_title>pt_expt: auto-generate forward/forward_lower in torch_module decorator to avoid boilerplate</issue_title> > <issue_description>## Summary > > The `torch_module` decorator in `deepmd/pt_expt/common.py` currently handles `__init__`, `__call__`, and `__setattr__` automatically, but each wrapped class still needs to manually define `forward()` (and potentially `forward_lower()`) methods that simply redirect to `call()` (and `call_lower()`). > > ## Current situation > > Every descriptor class in `deepmd/pt_expt/descriptor/` repeats the same boilerplate pattern: > > ```python > @torch_module > class DescrptSeA(DescrptSeADP): > def forward( > self, > extended_coord: torch.Tensor, > extended_atype: torch.Tensor, > nlist: torch.Tensor, > mapping: torch.Tensor | None = None, > ) -> tuple[...]: > descrpt, rot_mat, g2, h2, sw = self.call( > extended_coord, > extended_atype, > nlist, > mapping=mapping, > ) > return descrpt, rot_mat, g2, h2, sw > ``` > > This identical `forward → call` redirect is duplicated across `DescrptSeA`, `DescrptSeR`, `DescrptSeT`, `DescrptSeTTebd`, `DescrptBlockSeTTebd`, and similarly for fitting classes like `InvarFitting` (though `InvarFitting` does not use `torch_module`). > > ## Proposal > > Modify the `torch_module` decorator to **automatically generate `forward` and `forward_lower`** methods that delegate to `call` and `call_lower` respectively, if: > > 1. The wrapped class has a `call` method (from the dpmodel base) but does **not** define its own `forward`. > 2. Similarly for `call_lower` → `forward_lower`. > > This could be implemented by adding something like the following inside `torch_module`: > > ```python > def torch_module(module: type[NativeOP]) -> type[torch.nn.Module]: > @wraps(module, updated=()) > class TorchModule(module, torch.nn.Module): > def __init__(self, *args, **kwargs): > torch.nn.Module.__init__(self) > module.__init__(self, *args, **kwargs) > > def __call__(self, *args, **kwargs): > return torch.nn.Module.__call__(self, *args, **kwargs) > > def __setattr__(self, name, value): > handled, value = dpmodel_setattr(self, name, value) > if not handled: > super().__setattr__(name, value) > > # Auto-generate forward -> call redirect > if hasattr(module, "call") and "forward" not in module.__dict__: > def forward(self, *args, **kwargs): > return self.call(*args, **kwargs) > TorchModule.forward = forward > > # Auto-generate forward_lower -> call_lower redirect > if hasattr(module, "call_lower") and "forward_lower" not in module.__dict__: > def forward_lower(self, *args, **kwargs): > return self.call_lower(*args, **kwargs) > TorchModule.forward_lower = forward_lower > > return TorchModule > ``` > > ## Benefits > > - **Eliminates boilerplate**: Descriptor/fitting wrapper classes become minimal (just the `@torch_module` decorator + `@BaseDescriptor.register` + class declaration). > - **Reduces maintenance burden**: Adding new descriptors to `pt_expt` requires fewer lines and less copy-paste. > - **Consistent behavior**: The `forward ↔ call` contract is enforced in one place rather than scattered across many files. > - **Less error-prone**: No risk of forgetting to update `forward` when `call` signature changes. > > ## Example after this change > > ```python > @BaseDescriptor.register("se_e2_a") > @BaseDescriptor.register("se_a") > @torch_module > class DescrptSeA(DescrptSeADP): > pass # forward() auto-generated from call() > ``` > > --- > *Submitted by OpenClaw on behalf of @njzjz*</issue_description> > > ## Comments on the Issue (you are @copilot in this section) > > <comments> > </comments> > </details>  - Fixes #5245  --- 💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more [Copilot coding agent tips](https://gh.io/copilot-coding-agent-tips) in the docs. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py
@@ -300,6 +300,20 @@ def torch_module(
 ) -> type[torch.nn.Module]:
     """Convert a NativeOP to a torch.nn.Module.
 
+    This decorator wraps a NativeOP class to make it a PyTorch module, handling
+    initialization, attribute setting, and method delegation automatically.
+
+    **Auto-generated methods:**
+
+    - If the wrapped class has a ``call()`` method but does not explicitly define
+      ``forward()``, a ``forward()`` method will be auto-generated that delegates
+      to ``call()``.
+    - If the wrapped class has a ``call_lower()`` method but does not explicitly
+      define ``forward_lower()``, a ``forward_lower()`` method will be auto-generated
+      that delegates to ``call_lower()``.
+    - Explicit ``forward()`` or ``forward_lower()`` definitions in the wrapped class
+      are always respected and will not be overridden.
+
     Parameters
     ----------
     module : type[NativeOP]
@@ -308,13 +322,13 @@ def torch_module(
     Returns
     -------
     type[torch.nn.Module]
-        The torch.nn.Module.
+        The torch.nn.Module with auto-generated delegation methods if applicable.
 
     Examples
     --------
     >>> @torch_module
     ... class MyModule(NativeOP):
-    ...     pass
+    ...     pass  # forward() auto-generated from call() if it exists
     """
 
     @wraps(module, updated=())
@@ -332,6 +346,22 @@ def __setattr__(self, name: str, value: Any) -> None:
             if not handled:
                 super().__setattr__(name, value)
 
+    # Auto-generate forward -> call redirect if not explicitly defined
+    if hasattr(module, "call") and "forward" not in module.__dict__:
+
+        def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN001
+            return self.call(*args, **kwargs)
+
+        TorchModule.forward = forward
+
+    # Auto-generate forward_lower -> call_lower redirect if not explicitly defined
+    if hasattr(module, "call_lower") and "forward_lower" not in module.__dict__:
+
+        def forward_lower(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN001
+            return self.call_lower(*args, **kwargs)
+
+        TorchModule.forward_lower = forward_lower
+
     return TorchModule
 
 
diff --git a/deepmd/pt_expt/descriptor/se_e2_a.py b/deepmd/pt_expt/descriptor/se_e2_a.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeAArrayAPI as DescrptSeADP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -15,23 +13,4 @@
 @BaseDescriptor.register("se_a")
 @torch_module
 class DescrptSeA(DescrptSeADP):
-    def forward(
-        self,
-        extended_coord: torch.Tensor,
-        extended_atype: torch.Tensor,
-        nlist: torch.Tensor,
-        mapping: torch.Tensor | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        descrpt, rot_mat, g2, h2, sw = self.call(
-            extended_coord,
-            extended_atype,
-            nlist,
-            mapping=mapping,
-        )
-        return descrpt, rot_mat, g2, h2, sw
+    pass
diff --git a/deepmd/pt_expt/descriptor/se_r.py b/deepmd/pt_expt/descriptor/se_r.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.descriptor.se_r import DescrptSeR as DescrptSeRDP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -15,23 +13,4 @@
 @BaseDescriptor.register("se_r")
 @torch_module
 class DescrptSeR(DescrptSeRDP):
-    def forward(
-        self,
-        extended_coord: torch.Tensor,
-        extended_atype: torch.Tensor,
-        nlist: torch.Tensor,
-        mapping: torch.Tensor | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        descrpt, rot_mat, g2, h2, sw = self.call(
-            extended_coord,
-            extended_atype,
-            nlist,
-            mapping=mapping,
-        )
-        return descrpt, rot_mat, g2, h2, sw
+    pass
diff --git a/deepmd/pt_expt/descriptor/se_t.py b/deepmd/pt_expt/descriptor/se_t.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.descriptor.se_t import DescrptSeT as DescrptSeTDP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -16,23 +14,4 @@
 @BaseDescriptor.register("se_a_3be")
 @torch_module
 class DescrptSeT(DescrptSeTDP):
-    def forward(
-        self,
-        extended_coord: torch.Tensor,
-        extended_atype: torch.Tensor,
-        nlist: torch.Tensor,
-        mapping: torch.Tensor | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        descrpt, rot_mat, g2, h2, sw = self.call(
-            extended_coord,
-            extended_atype,
-            nlist,
-            mapping=mapping,
-        )
-        return descrpt, rot_mat, g2, h2, sw
+    pass
diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdDP
 from deepmd.pt_expt.common import (
     torch_module,
@@ -14,23 +12,4 @@
 @BaseDescriptor.register("se_e3_tebd")
 @torch_module
 class DescrptSeTTebd(DescrptSeTTebdDP):
-    def forward(
-        self,
-        extended_coord: torch.Tensor,
-        extended_atype: torch.Tensor,
-        nlist: torch.Tensor,
-        mapping: torch.Tensor | None = None,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-    ]:
-        descrpt, rot_mat, g2, h2, sw = self.call(
-            extended_coord,
-            extended_atype,
-            nlist,
-            mapping=mapping,
-        )
-        return descrpt, rot_mat, g2, h2, sw
+    pass
diff --git a/deepmd/pt_expt/descriptor/se_t_tebd_block.py b/deepmd/pt_expt/descriptor/se_t_tebd_block.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.descriptor.se_t_tebd import (
     DescrptBlockSeTTebd as DescrptBlockSeTTebdDP,
 )
@@ -13,29 +11,7 @@
 
 @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,
-        )
+    pass
 
 
 register_dpmodel_mapping(
diff --git a/deepmd/pt_expt/fitting/ener_fitting.py b/deepmd/pt_expt/fitting/ener_fitting.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
@@ -21,25 +19,7 @@ class EnergyFittingNet(EnergyFittingNetDP):
     This inherits from dpmodel EnergyFittingNet to get the correct serialize() method.
     """
 
-    def forward(
-        self,
-        descriptor: torch.Tensor,
-        atype: torch.Tensor,
-        gr: torch.Tensor | None = None,
-        g2: torch.Tensor | None = None,
-        h2: torch.Tensor | None = None,
-        fparam: torch.Tensor | None = None,
-        aparam: torch.Tensor | None = None,
-    ) -> dict[str, torch.Tensor]:
-        return self.call(
-            descriptor,
-            atype,
-            gr=gr,
-            g2=g2,
-            h2=h2,
-            fparam=fparam,
-            aparam=aparam,
-        )
+    pass
 
 
 register_dpmodel_mapping(
diff --git a/deepmd/pt_expt/fitting/invar_fitting.py b/deepmd/pt_expt/fitting/invar_fitting.py
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-import torch
-
 from deepmd.dpmodel.fitting.invar_fitting import InvarFitting as InvarFittingDP
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
@@ -15,25 +13,7 @@
 @BaseFitting.register("invar")
 @torch_module
 class InvarFitting(InvarFittingDP):
-    def forward(
-        self,
-        descriptor: torch.Tensor,
-        atype: torch.Tensor,
-        gr: torch.Tensor | None = None,
-        g2: torch.Tensor | None = None,
-        h2: torch.Tensor | None = None,
-        fparam: torch.Tensor | None = None,
-        aparam: torch.Tensor | None = None,
-    ) -> dict[str, torch.Tensor]:
-        return self.call(
-            descriptor,
-            atype,
-            gr=gr,
-            g2=g2,
-            h2=h2,
-            fparam=fparam,
-            aparam=aparam,
-        )
+    pass
 
 
 register_dpmodel_mapping(
diff --git a/deepmd/pt_expt/utils/network.py b/deepmd/pt_expt/utils/network.py
@@ -194,9 +194,6 @@ def __init__(self, layers: list[dict] | None = None) -> None:
         super().__init__(layers)
         self.layers = torch.nn.ModuleList(self.layers)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.call(x)
-
 
 class EmbeddingNet(EmbeddingNetDP, torch.nn.Module):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
diff --git a/source/tests/pt_expt/utils/test_common.py b/source/tests/pt_expt/utils/test_common.py
