feat(pt_expt): add fitting for energy (#5218)

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

* **Bug Fixes**
* Improved tensor/device and dtype handling for consistent behavior
across NumPy and PyTorch backends.
* Fixed deserialization when layer collections are empty to avoid
errors.

* **New Features**
* Added experimental PyTorch fitting wrappers: EnergyFittingNet and
InvarFitting for tensor-based workflows and export/tracing.
  * Renamed descriptor registration keys to streamlined identifiers.
  * Package exports updated to surface new fitting entry points.

* **Tests**
* Added extensive tests for energy/invariant fitting, statistics
computation, and PyTorch export compatibility.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/dpmodel/fitting/general_fitting.py

@@ -261,8 +261,18 @@ def compute_input_stats(
                 fparam_std,
             )
             fparam_inv_std = 1.0 / fparam_std
-            self.fparam_avg = fparam_avg.astype(self.fparam_avg.dtype)
-            self.fparam_inv_std = fparam_inv_std.astype(self.fparam_inv_std.dtype)
+            # Use array_api_compat to handle both numpy and torch
+            xp = array_api_compat.array_namespace(self.fparam_avg)
+            self.fparam_avg = xp.asarray(
+                fparam_avg,
+                dtype=self.fparam_avg.dtype,
+                device=array_api_compat.device(self.fparam_avg),
+            )
+            self.fparam_inv_std = xp.asarray(
+                fparam_inv_std,
+                dtype=self.fparam_inv_std.dtype,
+                device=array_api_compat.device(self.fparam_inv_std),
+            )
         # stat aparam
         if self.numb_aparam > 0:
             sys_sumv = []
@@ -284,8 +294,18 @@ def compute_input_stats(
                 aparam_std,
             )
             aparam_inv_std = 1.0 / aparam_std
-            self.aparam_avg = aparam_avg.astype(self.aparam_avg.dtype)
-            self.aparam_inv_std = aparam_inv_std.astype(self.aparam_inv_std.dtype)
+            # Use array_api_compat to handle both numpy and torch
+            xp = array_api_compat.array_namespace(self.aparam_avg)
+            self.aparam_avg = xp.asarray(
+                aparam_avg,
+                dtype=self.aparam_avg.dtype,
+                device=array_api_compat.device(self.aparam_avg),
+            )
+            self.aparam_inv_std = xp.asarray(
+                aparam_inv_std,
+                dtype=self.aparam_inv_std.dtype,
+                device=array_api_compat.device(self.aparam_inv_std),
+            )
 
     @abstractmethod
     def _net_out_dim(self) -> int:
@@ -566,7 +586,9 @@ def _call_common(
         # calculate the prediction
         if not self.mixed_types:
             outs = xp.zeros(
-                [nf, nloc, net_dim_out], dtype=get_xp_precision(xp, self.precision)
+                [nf, nloc, net_dim_out],
+                dtype=get_xp_precision(xp, self.precision),
+                device=array_api_compat.device(descriptor),
             )
             for type_i in range(self.ntypes):
                 mask = xp.tile(

deepmd/dpmodel/utils/network.py

@@ -1110,10 +1110,13 @@ def deserialize(cls, data: dict) -> "FittingNet":
         layers = data.pop("layers")
         obj = cls(**data)
         # Use type(obj.layers[0]) to respect subclass layer types
-        layer_type = type(obj.layers[0])
-        obj.layers = type(obj.layers)(
-            [layer_type.deserialize(layer) for layer in layers]
-        )
+        if obj.layers:
+            layer_type = type(obj.layers[0])
+            obj.layers = type(obj.layers)(
+                [layer_type.deserialize(layer) for layer in layers]
+            )
+        else:
+            obj.layers = type(obj.layers)([])
         return obj
 
 

deepmd/pt_expt/descriptor/se_e2_a.py

@@ -11,8 +11,8 @@
 )
 
 
-@BaseDescriptor.register("se_e2_a_expt")
-@BaseDescriptor.register("se_a_expt")
+@BaseDescriptor.register("se_e2_a")
+@BaseDescriptor.register("se_a")
 @torch_module
 class DescrptSeA(DescrptSeADP):
     def forward(

deepmd/pt_expt/descriptor/se_r.py

@@ -11,8 +11,8 @@
 )
 
 
-@BaseDescriptor.register("se_e2_r_expt")
-@BaseDescriptor.register("se_r_expt")
+@BaseDescriptor.register("se_e2_r")
+@BaseDescriptor.register("se_r")
 @torch_module
 class DescrptSeR(DescrptSeRDP):
     def forward(

deepmd/pt_expt/descriptor/se_t.py

@@ -11,9 +11,9 @@
 )
 
 
-@BaseDescriptor.register("se_e3_expt")
-@BaseDescriptor.register("se_at_expt")
-@BaseDescriptor.register("se_a_3be_expt")
+@BaseDescriptor.register("se_e3")
+@BaseDescriptor.register("se_at")
+@BaseDescriptor.register("se_a_3be")
 @torch_module
 class DescrptSeT(DescrptSeTDP):
     def forward(

deepmd/pt_expt/descriptor/se_t_tebd.py

@@ -11,7 +11,7 @@
 )
 
 
-@BaseDescriptor.register("se_e3_tebd_expt")
+@BaseDescriptor.register("se_e3_tebd")
 @torch_module
 class DescrptSeTTebd(DescrptSeTTebdDP):
     def forward(

deepmd/pt_expt/fitting/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .base_fitting import (
+    BaseFitting,
+)
+from .ener_fitting import (
+    EnergyFittingNet,
+)
+from .invar_fitting import (
+    InvarFitting,
+)
+
+__all__ = [
+    "BaseFitting",
+    "EnergyFittingNet",
+    "InvarFitting",
+]

deepmd/pt_expt/fitting/base_fitting.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import torch
+
+from deepmd.dpmodel.fitting import (
+    make_base_fitting,
+)
+
+BaseFitting = make_base_fitting(torch.Tensor, "forward")

deepmd/pt_expt/fitting/ener_fitting.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import torch
+
+from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP
+from deepmd.pt_expt.common import (
+    dpmodel_setattr,
+    register_dpmodel_mapping,
+)
+from deepmd.pt_expt.utils.network import (
+    NetworkCollection,
+)
+
+from .base_fitting import (
+    BaseFitting,
+)
+
+
+@BaseFitting.register("ener")
+class EnergyFittingNet(EnergyFittingNetDP, torch.nn.Module):
+    """Energy fitting net for pt_expt backend.
+
+    This inherits from dpmodel EnergyFittingNet to get the correct serialize() method.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        EnergyFittingNetDP.__init__(self, *args, **kwargs)
+        # Convert dpmodel NetworkCollection to pt_expt NetworkCollection
+        self.nets = NetworkCollection.deserialize(self.nets.serialize())
+
+    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)
+
+    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,
+        )
+
+
+register_dpmodel_mapping(
+    EnergyFittingNetDP,
+    lambda v: EnergyFittingNet.deserialize(v.serialize()),
+)

deepmd/pt_expt/fitting/invar_fitting.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+import torch
+
+from deepmd.dpmodel.fitting.invar_fitting import InvarFitting as InvarFittingDP
+from deepmd.pt_expt.common import (
+    dpmodel_setattr,
+    register_dpmodel_mapping,
+)
+from deepmd.pt_expt.fitting.base_fitting import (
+    BaseFitting,
+)
+from deepmd.pt_expt.utils.network import (
+    NetworkCollection,
+)
+
+
+@BaseFitting.register("invar")
+class InvarFitting(InvarFittingDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        InvarFittingDP.__init__(self, *args, **kwargs)
+        # Convert dpmodel NetworkCollection to pt_expt NetworkCollection
+        self.nets = NetworkCollection.deserialize(self.nets.serialize())
+
+    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)
+
+    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,
+        )
+
+
+register_dpmodel_mapping(
+    InvarFittingDP,
+    lambda v: InvarFitting.deserialize(v.serialize()),
+)