refact(dpmodel,pt_expt): fitting net (#5207)

# FittingNet Refactoring: Factory Function to Concrete Class

## Summary

This refactoring converts `FittingNet` from a factory-generated dynamic
class to a concrete class in the dpmodel backend, following the same
pattern as the EmbeddingNet refactoring. This enables the auto-detection
registry mechanism in pt_expt to work seamlessly with FittingNet.

This PR is considered after
#5194 and
#5204

## Motivation

**Before**: `FittingNet` was created by a factory function
`make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)`, producing
a dynamically-typed class. This caused:

1. **Cannot be registered**: Dynamic classes can't be imported or
registered at module import time in the pt_expt registry
2. **Type matching fails**: Each call to `make_fitting_network` creates
a new class type, so registry lookup by type fails

**After**: `FittingNet` is now a concrete class that can be registered
in the pt_expt auto-conversion registry.

## Changes

### 1. dpmodel: Concrete `FittingNet` class

**File**: `deepmd/dpmodel/utils/network.py`

- Created concrete `FittingNet(EmbeddingNet)` class
- Moved constructor logic from factory into `__init__`
- Fixed `deserialize` to use `type(obj.layers[0])` instead of hardcoding
`T_Network.__init__(obj, layers)`, allowing pt_expt subclass to preserve
its converted torch layers
- Kept `make_fitting_network` factory for backwards compatibility (for
pt/pd backends)

```python
class FittingNet(EmbeddingNet):
    """The fitting network."""

    def __init__(self, in_dim, out_dim, neuron=[24, 48, 96],
                 activation_function="tanh", resnet_dt=False,
                 precision=DEFAULT_PRECISION, bias_out=True,
                 seed=None, trainable=True):
        # Handle trainable parameter
        if trainable is None:
            trainable = [True] * (len(neuron) + 1)
        elif isinstance(trainable, bool):
            trainable = [trainable] * (len(neuron) + 1)

        # Initialize embedding layers via parent
        super().__init__(
            in_dim, neuron=neuron,
            activation_function=activation_function,
            resnet_dt=resnet_dt, precision=precision,
            seed=seed, trainable=trainable[:-1]
        )

        # Add output layer
        i_in = neuron[-1] if len(neuron) > 0 else in_dim
        self.layers.append(
            NativeLayer(
                i_in, out_dim, bias=bias_out,
                use_timestep=False, activation_function=None,
                resnet=False, precision=precision,
                seed=child_seed(seed, len(neuron)),
                trainable=trainable[-1]
            )
        )
        self.out_dim = out_dim
        self.bias_out = bias_out

    @classmethod
    def deserialize(cls, data):
        data = data.copy()
        check_version_compatibility(data.pop("@Version", 1), 1, 1)
        data.pop("@Class", None)
        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]
        )
        return obj
```

### 2. pt_expt: Wrapper and registration

**File**: `deepmd/pt_expt/utils/network.py`

- Added import: `from deepmd.dpmodel.utils.network import FittingNet as
FittingNetDP`
- Created `FittingNet(FittingNetDP, torch.nn.Module)` wrapper
- Converts dpmodel layers to pt_expt `NativeLayer` (torch modules) in
`__init__`
- Registered in auto-conversion registry

```python
from deepmd.dpmodel.utils.network import FittingNet as FittingNetDP

class FittingNet(FittingNetDP, torch.nn.Module):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        torch.nn.Module.__init__(self)
        FittingNetDP.__init__(self, *args, **kwargs)
        # Convert dpmodel layers to pt_expt NativeLayer
        self.layers = torch.nn.ModuleList(
            [NativeLayer.deserialize(layer.serialize()) for layer in 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)

register_dpmodel_mapping(
    FittingNetDP,
    lambda v: FittingNet.deserialize(v.serialize()),
)
```

## Tests

### dpmodel tests

**File**: `source/tests/common/dpmodel/test_network.py`

Added to `TestFittingNet` class:

1. **`test_fitting_net`**: Original roundtrip serialization test
(already existed)
2. **`test_is_concrete_class`**: Verifies `FittingNet` is now a concrete
class, not factory output
3. **`test_forward_pass`**: Tests dpmodel forward pass produces correct
output shapes (single and batch)
4. **`test_trainable_parameter_variants`**: Tests different trainable
configurations (all trainable, all frozen, mixed)

### pt_expt integration tests

**File**: `source/tests/pt_expt/utils/test_network.py`

Created `TestFittingNetRefactor` test suite with 4 tests:

1. **`test_pt_expt_fitting_net_wraps_dpmodel`**: Verifies pt_expt
wrapper inherits correctly and converts layers
2. **`test_pt_expt_fitting_net_forward`**: Tests pt_expt forward pass
returns torch.Tensor with correct shape
3. **`test_serialization_round_trip_pt_expt`**: Tests pt_expt
serialize/deserialize round-trip
4. **`test_registry_converts_dpmodel_to_pt_expt`**: Tests
`try_convert_module` auto-converts dpmodel to pt_expt

## Verification

All tests pass:

```bash
# dpmodel network tests (includes new FittingNet tests)
python -m pytest source/tests/common/dpmodel/test_network.py -v
# 19 passed in 0.56s (was 16, added 3 FittingNet tests)

# dpmodel FittingNet tests specifically
python -m pytest source/tests/common/dpmodel/test_network.py::TestFittingNet -v
# 4 passed in 0.44s

# pt_expt network tests (EmbeddingNet + FittingNet)
python -m pytest source/tests/pt_expt/utils/test_network.py -v
# 14 passed in 0.45s

# Descriptor tests (verify refactoring doesn't break existing code)
python -m pytest source/tests/pt_expt/descriptor/ -v
# 8 passed in 5.43s
```

## Benefits

1. **Type-based auto-detection**: FittingNet now works with the registry
mechanism
2. **Consistency**: Same pattern as EmbeddingNet and other dpmodel
classes
3. **Maintainability**: Single source of truth for FittingNet in dpmodel
4. **Future-proof**: Any dpmodel FittingNet instances can be
auto-converted to pt_expt

## Backward Compatibility

- Serialization format unchanged (version 1)
- All existing tests pass
- `make_fitting_network` factory kept for pt/pd backends
- No changes to public API

## Files Changed

### Modified

- `deepmd/dpmodel/utils/network.py`: Concrete FittingNet class +
deserialize fix
- `deepmd/pt_expt/utils/network.py`: FittingNet wrapper + registration
- `source/tests/common/dpmodel/test_network.py`: Added dpmodel
FittingNet tests (3 new tests)
- `source/tests/pt_expt/utils/test_network.py`: Added pt_expt
integration tests (4 new tests)

### Pattern

This refactoring follows the exact same pattern as
`EMBEDDING_NET_REFACTOR.md`:

1. Convert factory-generated class to concrete class in dpmodel
2. Fix `deserialize` to use `type(obj.layers[0])`
3. Create pt_expt wrapper with layer conversion in `__init__`
4. Register with `register_dpmodel_mapping`
5. Add comprehensive tests


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

## Release Notes

* **New Features**
* Added PyTorch experimental descriptor implementations for SeT and
SeTTebd with full export/tracing support
* Introduced PyTorch-compatible wrapper classes for network components
enabling seamless integration with PyTorch workflows

* **Improvements**
* Enhanced device-aware tensor operations across all descriptors for
better multi-device support
* Improved error handling with explicit error messages when statistics
are missing instead of silent failures
* Refactored FittingNet as a concrete class with explicit public
interface

* **Tests**
* Added comprehensive test coverage for new PyTorch experimental
descriptors and network wrappers
* Added unit tests validating serialization, deserialization, and
forward pass behavior
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

deepmd/dpmodel/utils/network.py

@@ -1003,7 +1003,118 @@ def deserialize(cls, data: dict) -> "FittingNet":
     return FN
 
 
-FittingNet = make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)
+class FittingNet(EmbeddingNet):
+    """The fitting network. It may be implemented as an embedding
+    net connected with a linear output layer.
+
+    Parameters
+    ----------
+    in_dim
+        Input dimension.
+    out_dim
+        Output dimension
+    neuron
+        The number of neurons in each hidden layer.
+    activation_function
+        The activation function.
+    resnet_dt
+        Use time step at the resnet architecture.
+    precision
+        Floating point precision for the model parameters.
+    bias_out
+        The last linear layer has bias.
+    seed : int, optional
+        Random seed.
+    trainable : bool or list[bool], optional
+        Whether the network is trainable.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        neuron: list[int] = [24, 48, 96],
+        activation_function: str = "tanh",
+        resnet_dt: bool = False,
+        precision: str = DEFAULT_PRECISION,
+        bias_out: bool = True,
+        seed: int | list[int] | None = None,
+        trainable: bool | list[bool] = True,
+    ) -> None:
+        if trainable is None:
+            trainable = [True] * (len(neuron) + 1)
+        elif isinstance(trainable, bool):
+            trainable = [trainable] * (len(neuron) + 1)
+        else:
+            pass
+        super().__init__(
+            in_dim,
+            neuron=neuron,
+            activation_function=activation_function,
+            resnet_dt=resnet_dt,
+            precision=precision,
+            seed=seed,
+            trainable=trainable[:-1],
+        )
+        i_in = neuron[-1] if len(neuron) > 0 else in_dim
+        i_ot = out_dim
+        self.layers.append(
+            NativeLayer(
+                i_in,
+                i_ot,
+                bias=bias_out,
+                use_timestep=False,
+                activation_function=None,
+                resnet=False,
+                precision=precision,
+                seed=child_seed(seed, len(neuron)),
+                trainable=trainable[-1],
+            )
+        )
+        self.out_dim = out_dim
+        self.bias_out = bias_out
+
+    def serialize(self) -> dict:
+        """Serialize the network to a dict.
+
+        Returns
+        -------
+        dict
+            The serialized network.
+        """
+        return {
+            "@class": "FittingNetwork",
+            "@version": 1,
+            "in_dim": self.in_dim,
+            "out_dim": self.out_dim,
+            "neuron": self.neuron.copy(),
+            "activation_function": self.activation_function,
+            "resnet_dt": self.resnet_dt,
+            "precision": self.precision,
+            "bias_out": self.bias_out,
+            "layers": [layer.serialize() for layer in self.layers],
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "FittingNet":
+        """Deserialize the network from a dict.
+
+        Parameters
+        ----------
+        data : dict
+            The dict to deserialize from.
+        """
+        data = data.copy()
+        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        data.pop("@class", None)
+        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]
+        )
+        return obj
 
 
 class NetworkCollection:

deepmd/pt_expt/utils/network.py

@@ -11,11 +11,11 @@
     NativeOP,
 )
 from deepmd.dpmodel.utils.network import EmbeddingNet as EmbeddingNetDP
+from deepmd.dpmodel.utils.network import FittingNet as FittingNetDP
 from deepmd.dpmodel.utils.network import LayerNorm as LayerNormDP
 from deepmd.dpmodel.utils.network import NativeLayer as NativeLayerDP
 from deepmd.dpmodel.utils.network import NetworkCollection as NetworkCollectionDP
 from deepmd.dpmodel.utils.network import (
-    make_fitting_network,
     make_multilayer_network,
 )
 from deepmd.pt_expt.common import (
@@ -114,8 +114,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 )
 
 
-class FittingNet(make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)):
-    pass
+class FittingNet(FittingNetDP, torch.nn.Module):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        torch.nn.Module.__init__(self)
+        FittingNetDP.__init__(self, *args, **kwargs)
+        # Convert dpmodel layers to pt_expt NativeLayer
+        self.layers = torch.nn.ModuleList(
+            [NativeLayer.deserialize(layer.serialize()) for layer in 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)
+
+
+register_dpmodel_mapping(
+    FittingNetDP,
+    lambda v: FittingNet.deserialize(v.serialize()),
+)
 
 
 class NetworkCollection(NetworkCollectionDP, torch.nn.Module):

source/tests/common/dpmodel/test_network.py

@@ -313,6 +313,104 @@ def test_fitting_net(self) -> None:
             en1.call(inp)
             np.testing.assert_allclose(en0.call(inp), en1.call(inp))
 
+    def test_is_concrete_class(self) -> None:
+        """Verify FittingNet is a concrete class, not factory-generated."""
+        in_dim = 4
+        out_dim = 1
+        neuron = [8, 16]
+        net = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+            bias_out=True,
+        )
+        # Check it's the actual FittingNet class, not a dynamic class
+        self.assertEqual(net.__class__.__name__, "FittingNet")
+        self.assertEqual(net.__class__.__module__, "deepmd.dpmodel.utils.network")
+        # Verify it has the expected attributes
+        self.assertEqual(net.in_dim, in_dim)
+        self.assertEqual(net.out_dim, out_dim)
+        self.assertEqual(net.neuron, neuron)
+        self.assertEqual(net.activation_function, "tanh")
+        self.assertEqual(net.resnet_dt, True)
+        self.assertEqual(net.bias_out, True)
+        # FittingNet has len(neuron) embedding layers + 1 output layer
+        self.assertEqual(len(net.layers), len(neuron) + 1)
+
+    def test_forward_pass(self) -> None:
+        """Test FittingNet forward pass produces correct output shape."""
+        in_dim = 4
+        out_dim = 3
+        neuron = [8, 16, 32]
+        net = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            activation_function="tanh",
+            resnet_dt=True,
+            precision="float64",
+        )
+        # Single sample
+        rng = np.random.default_rng()
+        x = rng.standard_normal(in_dim)
+        out = net.call(x)
+        self.assertEqual(out.shape, (out_dim,))
+
+        # Batch of samples
+        batch_size = 5
+        x_batch = rng.standard_normal((batch_size, in_dim))
+        out_batch = net.call(x_batch)
+        self.assertEqual(out_batch.shape, (batch_size, out_dim))
+
+    def test_trainable_parameter_variants(self) -> None:
+        """Test FittingNet with different trainable configurations."""
+        in_dim = 4
+        out_dim = 2
+        neuron = [8, 16]
+
+        # Test 1: All layers trainable (default)
+        net_all_trainable = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=True,
+        )
+        for layer in net_all_trainable.layers:
+            self.assertTrue(layer.trainable)
+
+        # Test 2: All layers frozen
+        net_all_frozen = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=False,
+        )
+        for layer in net_all_frozen.layers:
+            self.assertFalse(layer.trainable)
+
+        # Test 3: Mixed trainable (embedding layers frozen, output layer trainable)
+        trainable_list = [False, False, True]  # 2 embedding layers + 1 output layer
+        net_mixed = FittingNet(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            neuron=neuron,
+            trainable=trainable_list,
+        )
+        self.assertFalse(net_mixed.layers[0].trainable)  # First embedding layer
+        self.assertFalse(net_mixed.layers[1].trainable)  # Second embedding layer
+        self.assertTrue(net_mixed.layers[2].trainable)  # Output layer
+
+        # Test 4: Serialize/deserialize preserves trainable
+        serialized = net_mixed.serialize()
+        net_restored = FittingNet.deserialize(serialized)
+        for orig_layer, restored_layer in zip(
+            net_mixed.layers, net_restored.layers, strict=True
+        ):
+            self.assertEqual(orig_layer.trainable, restored_layer.trainable)
+
 
 class TestNetworkCollection(unittest.TestCase):
     def setUp(self) -> None:

source/tests/pt_expt/utils/test_network.py

@@ -281,3 +281,124 @@ def test_trainable_parameter_handling(self) -> None:
         for layer in net_frozen.layers:
             if layer.w is not None:
                 self.assertFalse(layer.w.requires_grad)
+
+
+class TestFittingNetRefactor(unittest.TestCase):
+    """Tests for the refactored FittingNet pt_expt wrapper."""
+
+    def setUp(self) -> None:
+        self.in_dim = 4
+        self.out_dim = 1
+        self.neuron = [8, 16]
+        self.activation = "tanh"
+        self.resnet_dt = True
+        self.precision = "float64"
+
+    def test_pt_expt_fitting_net_wraps_dpmodel(self) -> None:
+        """Verify pt_expt FittingNet correctly wraps dpmodel."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        # Check it's a torch.nn.Module
+        self.assertIsInstance(net, torch.nn.Module)
+        # Check layers are converted to pt_expt NativeLayer (torch modules)
+        self.assertIsInstance(net.layers, torch.nn.ModuleList)
+        for layer in net.layers:
+            self.assertIsInstance(layer, torch.nn.Module)
+
+    def test_pt_expt_fitting_net_forward(self) -> None:
+        """Test pt_expt FittingNet forward pass returns torch.Tensor."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        x = torch.randn(5, self.in_dim, dtype=torch.float64, device=env.DEVICE)
+        out = net(x)
+        self.assertIsInstance(out, torch.Tensor)
+        self.assertEqual(out.shape, (5, self.out_dim))
+        self.assertEqual(out.dtype, torch.float64)
+
+    def test_serialization_round_trip_pt_expt(self) -> None:
+        """Test pt_expt FittingNet serialization/deserialization."""
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        net = FittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+        x = torch.randn(5, self.in_dim, dtype=torch.float64, device=env.DEVICE)
+        out1 = net(x)
+
+        # Serialize and deserialize
+        serialized = net.serialize()
+        net2 = FittingNet.deserialize(serialized)
+
+        # Verify layers are still pt_expt NativeLayer modules
+        self.assertIsInstance(net2.layers, torch.nn.ModuleList)
+        for layer in net2.layers:
+            self.assertIsInstance(layer, torch.nn.Module)
+
+        out2 = net2(x)
+        np.testing.assert_allclose(
+            out1.detach().cpu().numpy(),
+            out2.detach().cpu().numpy(),
+        )
+
+    def test_registry_converts_dpmodel_to_pt_expt(self) -> None:
+        """Test that dpmodel FittingNet can be converted to pt_expt via registry."""
+        from deepmd.dpmodel.utils.network import FittingNet as DPFittingNet
+        from deepmd.pt_expt.common import (
+            try_convert_module,
+        )
+        from deepmd.pt_expt.utils.network import (
+            FittingNet,
+        )
+
+        # Create dpmodel FittingNet
+        dp_net = DPFittingNet(
+            in_dim=self.in_dim,
+            out_dim=self.out_dim,
+            neuron=self.neuron,
+            activation_function=self.activation,
+            resnet_dt=self.resnet_dt,
+            precision=self.precision,
+            seed=GLOBAL_SEED,
+        )
+
+        # Try to convert via registry
+        converted = try_convert_module(dp_net)
+
+        # Should return pt_expt FittingNet
+        self.assertIsNotNone(converted)
+        self.assertIsInstance(converted, torch.nn.Module)
+        self.assertIsInstance(converted, FittingNet)
+
+        # Verify layers are pt_expt modules
+        for layer in converted.layers:
+            self.assertIsInstance(layer, torch.nn.Module)