refact (pt_expt): provide infrastructure for converting dpmodel classes to PyTorch modules. by wanghan-iapcm · Pull Request #5204 · deepmodeling/deepmd-kit

wanghan-iapcm · 2026-02-08T13:21:49Z

consider after the merge of #5194

automatically wrapping dpmodel classes (array_api_compat-based) as PyTorch modules. The key insight is to detect attributes by their value type rather than by hard-coded names.

Summary by CodeRabbit

New Features
- Registry-driven conversion for DP objects to PyTorch modules enabling automatic wrapper creation.
- New PyTorch-friendly descriptor variants with stable forward outputs for se_e2_a and se_r.
- PyTorch-wrapped exclude-mask utilities and a NetworkCollection of wrapped network types for proper module/state handling.
- Device-aware tensor conversion and robust handling of numpy buffers and None-valued buffers for reliable serialization/movement.

…y on pt backend.

codecov · 2026-02-08T13:54:54Z

Codecov Report

❌ Patch coverage is 94.73684% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.00%. Comparing base (5c2ca51) to head (55e094e).
⚠️ Report is 159 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pt_expt/common.py	91.42%	3 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5204   +/-   ##
=======================================
  Coverage   81.99%   82.00%           
=======================================
  Files         724      724           
  Lines       73807    73801    -6     
  Branches     3616     3615    -1     
=======================================
+ Hits        60519    60520    +1     
+ Misses      12124    12118    -6     
+ Partials     1164     1163    -1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@Version

# EmbeddingNet Refactoring: Factory Function to Concrete Class ## Summary This refactoring converts `EmbeddingNet` from a factory-generated dynamic class to a concrete class in the dpmodel backend. This change enables the auto-detection registry mechanism in pt_expt to work seamlessly with EmbeddingNet attributes. This PR is considered after #5194 and #5204 ## Motivation **Before**: `EmbeddingNet` was created by a factory function `make_embedding_network(NativeNet, NativeLayer)`, producing a dynamically-typed class `make_embedding_network.<locals>.EN`. This caused two problems: 1. **Cannot be registered**: Dynamic classes can't be imported or registered at module import time in the pt_expt registry 2. **Name-based hacks required**: pt_expt wrappers had to explicitly check for `name == "embedding_net"` in `__setattr__` instead of using the type-based auto-detection mechanism **After**: `EmbeddingNet` is now a concrete class that can be registered in the pt_expt auto-conversion registry, eliminating the need for name-based special cases. ## Changes ### 1. dpmodel: Concrete `EmbeddingNet` class **File**: `deepmd/dpmodel/utils/network.py` - Replaced factory-generated class with concrete `EmbeddingNet(NativeNet)` class - Moved constructor logic from factory into `__init__` - Fixed `deserialize` to use `type(obj.layers[0])` instead of hardcoding `super(EmbeddingNet, obj)`, allowing pt_expt subclass to preserve its converted torch layers - Kept `make_embedding_network` factory for pt/pd backends that use different base classes (MLP) ```python class EmbeddingNet(NativeNet): """The embedding network.""" def __init__(self, in_dim, neuron=[24, 48, 96], activation_function="tanh", resnet_dt=False, precision=DEFAULT_PRECISION, seed=None, bias=True, trainable=True): layers = [] i_in = in_dim if isinstance(trainable, bool): trainable = [trainable] * len(neuron) for idx, ii in enumerate(neuron): i_ot = ii layers.append( NativeLayer( i_in, i_ot, bias=bias, use_timestep=resnet_dt, activation_function=activation_function, resnet=True, precision=precision, seed=child_seed(seed, idx), trainable=trainable[idx] ).serialize() ) i_in = i_ot super().__init__(layers) self.in_dim = in_dim self.neuron = neuron self.activation_function = activation_function self.resnet_dt = resnet_dt self.precision = precision self.bias = bias @classmethod def deserialize(cls, data): data = data.copy() check_version_compatibility(data.pop("@Version", 1), 2, 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` - Created `EmbeddingNet(EmbeddingNetDP, torch.nn.Module)` wrapper - Converts dpmodel layers to pt_expt `NativeLayer` (torch modules) in `__init__` - Registered in auto-conversion registry ```python class EmbeddingNet(EmbeddingNetDP, torch.nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: torch.nn.Module.__init__(self) EmbeddingNetDP.__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( EmbeddingNetDP, lambda v: EmbeddingNet.deserialize(v.serialize()), ) ``` ### 3. TypeEmbedNet: Simplified to use registry **File**: `deepmd/pt_expt/utils/type_embed.py` - No longer needs name-based `embedding_net` check in `__setattr__` - Uses common `dpmodel_setattr` which auto-converts via registry - Imports `network` module to ensure `EmbeddingNet` registration happens first ```python class TypeEmbedNet(TypeEmbedNetDP, torch.nn.Module): def __setattr__(self, name: str, value: Any) -> None: # Auto-converts embedding_net via registry handled, value = dpmodel_setattr(self, name, value) if not handled: super().__setattr__(name, value) ``` ## Tests ### dpmodel tests **File**: `source/tests/common/dpmodel/test_network.py` Added to `TestEmbeddingNet` class: 1. **`test_is_concrete_class`**: Verifies `EmbeddingNet` is now a concrete class, not factory output 2. **`test_forward_pass`**: Tests dpmodel forward pass produces correct shapes 3. **`test_trainable_parameter_variants`**: Tests different trainable configurations (all trainable, all frozen, mixed) (The existing `test_embedding_net` test already covers serialization/deserialization round-trip) ### pt_expt integration tests **File**: `source/tests/pt_expt/utils/test_network.py` Created `TestEmbeddingNetRefactor` test suite with 8 tests: 1. **`test_pt_expt_embedding_net_wraps_dpmodel`**: Verifies pt_expt wrapper inherits correctly and converts layers 2. **`test_pt_expt_embedding_net_forward`**: Tests pt_expt forward pass returns torch.Tensor 3. **`test_serialization_round_trip_pt_expt`**: Tests pt_expt serialize/deserialize 4. **`test_deserialize_preserves_layer_type`**: Tests the key fix - `deserialize` uses `type(obj.layers[0])` to preserve pt_expt's torch layers 5. **`test_cross_backend_consistency`**: Tests numerical consistency between dpmodel and pt_expt 6. **`test_registry_converts_dpmodel_to_pt_expt`**: Tests `try_convert_module` auto-converts dpmodel to pt_expt 7. **`test_auto_conversion_in_setattr`**: Tests `dpmodel_setattr` auto-converts EmbeddingNet attributes 8. **`test_trainable_parameter_handling`**: Tests trainable vs frozen parameters work correctly in pt_expt ## Verification All tests pass: ```bash # dpmodel EmbeddingNet tests python -m pytest source/tests/common/dpmodel/test_network.py::TestEmbeddingNet -v # 4 passed in 0.41s # pt_expt EmbeddingNet integration tests python -m pytest source/tests/pt_expt/utils/test_network.py::TestEmbeddingNetRefactor -v # 8 passed in 0.41s # All pt_expt network tests python -m pytest source/tests/pt_expt/utils/test_network.py -v # 10 passed in 0.41s # Descriptor tests (verify refactoring doesn't break existing code) python -m pytest source/tests/pt_expt/descriptor/test_se_e2_a.py -v -k consistency # 1 passed python -m pytest source/tests/universal/pt_expt/descriptor/test_descriptor.py -v # 8 passed in 3.27s ``` ## Benefits 1. **Type-based auto-detection**: No more name-based special cases in `__setattr__` 2. **Maintainability**: Single source of truth for EmbeddingNet in dpmodel 3. **Consistency**: Same pattern as other dpmodel classes (AtomExcludeMask, NetworkCollection, etc.) 4. **Future-proof**: New attributes in dpmodel automatically work in pt_expt via registry ## Backward Compatibility - Serialization format unchanged (version 2.1) - All existing tests pass - `make_embedding_network` factory kept for pt/pd backends - No changes to public API ## Files Changed ### Modified - `deepmd/dpmodel/utils/network.py`: Concrete EmbeddingNet class + deserialize fix - `deepmd/pt_expt/utils/network.py`: EmbeddingNet wrapper + registration - `deepmd/pt_expt/utils/type_embed.py`: Simplified to use registry - `source/tests/common/dpmodel/test_network.py`: Added dpmodel EmbeddingNet tests (3 new tests) - `source/tests/pt_expt/utils/test_network.py`: Added pt_expt integration tests (8 new tests) ### No changes required - All descriptor wrappers (se_e2_a, se_r, se_t, se_t_tebd) automatically work via registry - No changes to dpmodel logic or array_api_compat code  ## Summary by CodeRabbit ## Release Notes * **New Features** * Added PyTorch compatibility layer enabling DPModel neural network components to be used with PyTorch workflows for training and inference * Enhanced embedding network with explicit serialization and deserialization capabilities * **Refactor** * Restructured embedding network with explicit class design for improved type stability and control flow management  --------- 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>

@Version

# 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  ## 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  --------- 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>

This PR is considered after #5194 #5204 and #5205  ## Summary by CodeRabbit ## Release Notes * **New Features** * Added experimental PyTorch support for SeT and SeT-TEBD descriptors, enabling model training and serialization/export. * Introduced TypeEmbedNet wrapper for type embedding integration in PyTorch workflows. * **Bug Fixes** * Improved backend compatibility and device-aware tensor allocation across descriptor implementations. * Fixed PyTorch tensor indexing compatibility issues. * **Tests** * Added comprehensive test coverage for new experimental descriptors and consistency validation.  --------- 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>

Han Wang added 30 commits February 6, 2026 07:48

implement pytorch-exportable for se_e2_a descriptor

ec2e031

better type for xp.zeros

b8a48ff

implement env, base_descriptor and exclude_mask, remove the dependenc…

1cc001f

…y on pt backend.

mv to_torch_tensor to common

f2fbe88

simplify __init__ of the NaiveLayer

e2afbe9

fix bug

4ba511a

fix bug

fb9598a

simplify init method of se_e2_a descriptor. fig bug in consistent UT

fa03351

restructure the test folders. add test_common.

09b33f1

add test_exclusion_mask.py

67f2e54

fix poitential import issue in test.

f7d83dd

correct __call__(). fix bug

0c96bb6

fix registration issue

9dca912

fix pt-expt file extension

17f0a5d

fix(pt): expansion of get_default_nthreads()

8ce93ba

fix bug of intra-inter

3091988

fix bug of default dp inter value

85f0583

fix cicd

d33324d

feat: add support for se_r

4de9a56

fix device of xp array

f4dc0af

fix device of xp array

2384835

revert extend_coord_with_ghosts

9646d71

raise error for non-implemented methods

f270069

restore import torch

57433d3

fix(pt,pt-expt): guard thread setters

eedcbaf

make exclusion mask modules

d8b2cf4

fix(pt-expt): clear params on None

aeef15a

fix bug

8bdb1f8

utility to handel dpmodel -> pt_expt conversion

d3b01da

fix to_numpy_array device

3452a2a

wanghan-iapcm mentioned this pull request Feb 8, 2026

refact(dpmodel,pt_expt): embedding net #5205

Merged

This was referenced Feb 8, 2026

feat(pt_expt): implement se_t and se_t_tebd descriptors. #5206

Closed

refact(dpmodel,pt_expt): fitting net #5207

Merged

wanghan-iapcm changed the title ~~refact: provide infrastructure for converting dpmodel classes to PyTorch modules.~~ refact (pt_expt): provide infrastructure for converting dpmodel classes to PyTorch modules. Feb 8, 2026

njzjz reviewed Feb 8, 2026

View reviewed changes

Comment thread deepmd/pt_expt/common.py Outdated

njzjz reviewed Feb 8, 2026

View reviewed changes

Comment thread deepmd/pt_expt/common.py Outdated

merge master

de8f156

wanghan-iapcm mentioned this pull request Feb 8, 2026

feat(pt_expt): implement se_t and se_t_tebd descriptors. #5208

Merged

wanghan-iapcm added the Test CUDA Trigger test CUDA workflow label Feb 8, 2026

github-actions Bot removed the Test CUDA Trigger test CUDA workflow label Feb 8, 2026

better type checking

87e9b9d

wanghan-iapcm requested a review from njzjz February 9, 2026 11:54

Han Wang added 2 commits February 9, 2026 19:58

fix

ef84c6c

raise error

55e094e

njzjz approved these changes Feb 9, 2026

View reviewed changes

njzjz added this pull request to the merge queue Feb 9, 2026

Merged via the queue into deepmodeling:master with commit 97d8ded Feb 9, 2026
70 checks passed

coderabbitai Bot mentioned this pull request Feb 9, 2026

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

Merged

wanghan-iapcm deleted the refact-auto-setattr branch February 10, 2026 02:40

This was referenced Feb 12, 2026

feat(pt_expt): atomic model #5219

Closed

feat(pt_expt): full model and refact the module output names of dpmodel backend #5243

Closed

feat(pt_expt): add descriptors dpa1 dpa2 dpa3 and hybrid #5248

Merged

coderabbitai Bot mentioned this pull request Mar 17, 2026

feat(pt_expt): add dp compress support for pt_expt backend #5323

Merged

3 tasks

This was referenced Jun 13, 2026

refactor(jax): auto-convert dpmodel modules #5527

Merged

refactor(tests): auto-convert array-api-strict modules #5528

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refact (pt_expt): provide infrastructure for converting dpmodel classes to PyTorch modules. #5204

refact (pt_expt): provide infrastructure for converting dpmodel classes to PyTorch modules. #5204
njzjz merged 34 commits into
deepmodeling:masterfrom
wanghan-iapcm:refact-auto-setattr

wanghan-iapcm commented Feb 8, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

codecov Bot commented Feb 8, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

wanghan-iapcm commented Feb 8, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

codecov Bot commented Feb 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wanghan-iapcm commented Feb 8, 2026 •

edited by coderabbitai Bot

Loading

codecov Bot commented Feb 8, 2026 •

edited

Loading