Skip to content

Commit c381452

Browse files
wanghan-iapcmHan Wang
wanghan-iapcm
and
Han Wang
authored
fix(jax,tf2): deep-copy and propagate type_map in ZBL model factories (deepmodeling#5725)
## Problem Fixes deepmodeling#5676. The JAX and TF2 ZBL model factories (`get_zbl_model`), unlike the standard-model factories and the dpmodel ZBL path, mutated the caller's config dict in place (popping the descriptor and `fitting_net` `type`) and did not inject `type_map` into the descriptor and fitting sub-configs. The in-place mutation is inconsistent with the neighboring factories, and the missing `type_map` leaves the descriptor and fitting without the model type map that standard/dpmodel construction provides. ## Fix Deep-copy the input and inject `type_map` into the descriptor and fitting configs in both `get_zbl_model` factories, mirroring the dpmodel ZBL path. The issue also suspected the retained JAX checkpoint metadata could be corrupted by the in-place mutation. That does not actually happen: the JAX trainer passes a `deepcopy` of the model params into the factory, so the stored `model_def_script` is never mutated. No change is made for that sub-claim. ## Test Adds `source/tests/jax/test_zbl_model.py` and a TF2-gated `source/tests/consistent/test_tf2_zbl_model.py`, each asserting the factory leaves the input dict unchanged and that the constructed descriptor and fitting carry the model `type_map`. Both assertions fail on master (the input is mutated and the sub-config `type_map` is `None`) and pass with the fix. The TF2 test is gated on `INSTALLED_TF2` (`DEEPMD_TEST_TF2=1`). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved ZBL model setup so input configuration data is no longer modified during model creation. * Ensured type mappings are consistently applied to both descriptor and fitting settings when building ZBL models. * **Tests** * Added coverage for JAX and TF2 ZBL model factories to verify configuration immutability and correct type mapping behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent 87c7da8 commit c381452

4 files changed

Lines changed: 160 additions & 0 deletions

File tree

‎deepmd/jax/model/model.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,13 @@ def get_standard_model(data: dict) -> BaseModel:
6868

6969

7070
def get_zbl_model(data: dict) -> DPZBLModel:
71+
data = deepcopy(data)
7172
data["descriptor"]["ntypes"] = len(data["type_map"])
73+
data["descriptor"]["type_map"] = data["type_map"]
7274
descriptor_type = data["descriptor"].pop("type")
7375
descriptor = BaseDescriptor.get_class_by_type(descriptor_type)(**data["descriptor"])
7476
fitting_type = data["fitting_net"].pop("type")
77+
data["fitting_net"]["type_map"] = data["type_map"]
7578
if fitting_type == "ener":
7679
fitting = EnergyFittingNet(
7780
ntypes=descriptor.get_ntypes(),

‎deepmd/tf2/model/model.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,10 +66,13 @@ def get_standard_model(data: dict) -> BaseModel:
6666

6767

6868
def get_zbl_model(data: dict) -> DPZBLModel:
69+
data = deepcopy(data)
6970
data["descriptor"]["ntypes"] = len(data["type_map"])
71+
data["descriptor"]["type_map"] = data["type_map"]
7072
descriptor_type = data["descriptor"].pop("type")
7173
descriptor = BaseDescriptor.get_class_by_type(descriptor_type)(**data["descriptor"])
7274
fitting_type = data["fitting_net"].pop("type")
75+
data["fitting_net"]["type_map"] = data["type_map"]
7376
if fitting_type == "ener":
7477
fitting = EnergyFittingNet(
7578
ntypes=descriptor.get_ntypes(),

‎source/tests/consistent/test_tf2_zbl_model.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Test the TF2 ZBL model factory deep-copies input and injects type_map.
3+
4+
The TF2 (and JAX) ZBL factory used to mutate the caller's config in place and
5+
did not propagate ``type_map`` into the descriptor and fitting sub-configs,
6+
unlike the standard and dpmodel factories. This checks the factory leaves the
7+
input dict unchanged and that the constructed descriptor and fitting carry the
8+
model ``type_map``. Gated on the TF2 backend (``DEEPMD_TEST_TF2=1``).
9+
"""
10+
11+
import os
12+
import unittest
13+
from copy import (
14+
deepcopy,
15+
)
16+
17+
from .common import (
18+
INSTALLED_TF2,
19+
)
20+
21+
if INSTALLED_TF2:
22+
from deepmd.tf2.model.model import (
23+
get_zbl_model,
24+
)
25+
26+
TESTS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
27+
SRTAB = os.path.join(
28+
TESTS_DIR, "pt", "water", "data", "zbl_tab_potential", "H2O_tab_potential.txt"
29+
)
30+
31+
32+
def _zbl_config() -> dict:
33+
return {
34+
"type_map": ["O", "H", "B"],
35+
"use_srtab": SRTAB,
36+
"sw_rmin": 0.2,
37+
"sw_rmax": 4.0,
38+
"smin_alpha": 0.1,
39+
# ZBL wraps a linear atomic model, which requires a mixed-type descriptor
40+
"descriptor": {
41+
"type": "se_atten",
42+
"sel": 40,
43+
"rcut_smth": 0.5,
44+
"rcut": 4.0,
45+
"neuron": [3, 6],
46+
"axis_neuron": 2,
47+
"attn": 8,
48+
"attn_layer": 2,
49+
"attn_dotr": True,
50+
"attn_mask": False,
51+
"set_davg_zero": True,
52+
"type_one_side": True,
53+
"seed": 1,
54+
},
55+
"fitting_net": {
56+
"type": "ener",
57+
"neuron": [5, 5],
58+
"seed": 1,
59+
},
60+
}
61+
62+
63+
@unittest.skipUnless(INSTALLED_TF2, "TF2 backend is not installed")
64+
class TestTF2ZBLModelFactory(unittest.TestCase):
65+
def test_does_not_mutate_input(self) -> None:
66+
data = _zbl_config()
67+
orig = deepcopy(data)
68+
get_zbl_model(data)
69+
self.assertEqual(data, orig)
70+
71+
def test_injects_type_map_into_subconfigs(self) -> None:
72+
data = _zbl_config()
73+
model = get_zbl_model(data)
74+
dp_atomic = model.atomic_model.models[0]
75+
self.assertEqual(list(dp_atomic.descriptor.get_type_map()), data["type_map"])
76+
self.assertEqual(list(dp_atomic.fitting_net.get_type_map()), data["type_map"])
77+
78+
79+
if __name__ == "__main__":
80+
unittest.main()

‎source/tests/jax/test_zbl_model.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Test the JAX ZBL model factory deep-copies input and injects type_map.
3+
4+
The JAX (and TF2) ZBL factory used to mutate the caller's config in place (via
5+
``pop("type")``) and did not propagate ``type_map`` into the descriptor and
6+
fitting sub-configs, unlike the standard and dpmodel factories. This checks the
7+
factory leaves the input dict unchanged and that the constructed descriptor and
8+
fitting carry the model ``type_map``.
9+
"""
10+
11+
import os
12+
import unittest
13+
from copy import (
14+
deepcopy,
15+
)
16+
17+
from deepmd.jax.model.model import (
18+
get_zbl_model,
19+
)
20+
21+
TESTS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
22+
SRTAB = os.path.join(
23+
TESTS_DIR, "pt", "water", "data", "zbl_tab_potential", "H2O_tab_potential.txt"
24+
)
25+
26+
27+
def _zbl_config() -> dict:
28+
return {
29+
"type_map": ["O", "H", "B"],
30+
"use_srtab": SRTAB,
31+
"sw_rmin": 0.2,
32+
"sw_rmax": 4.0,
33+
"smin_alpha": 0.1,
34+
# ZBL wraps a linear atomic model, which requires a mixed-type descriptor
35+
"descriptor": {
36+
"type": "se_atten",
37+
"sel": 40,
38+
"rcut_smth": 0.5,
39+
"rcut": 4.0,
40+
"neuron": [3, 6],
41+
"axis_neuron": 2,
42+
"attn": 8,
43+
"attn_layer": 2,
44+
"attn_dotr": True,
45+
"attn_mask": False,
46+
"set_davg_zero": True,
47+
"type_one_side": True,
48+
"seed": 1,
49+
},
50+
"fitting_net": {
51+
"type": "ener",
52+
"neuron": [5, 5],
53+
"seed": 1,
54+
},
55+
}
56+
57+
58+
class TestJAXZBLModelFactory(unittest.TestCase):
59+
def test_does_not_mutate_input(self) -> None:
60+
data = _zbl_config()
61+
orig = deepcopy(data)
62+
get_zbl_model(data)
63+
self.assertEqual(data, orig)
64+
65+
def test_injects_type_map_into_subconfigs(self) -> None:
66+
data = _zbl_config()
67+
model = get_zbl_model(data)
68+
dp_atomic = model.atomic_model.models[0]
69+
self.assertEqual(dp_atomic.descriptor.get_type_map(), data["type_map"])
70+
self.assertEqual(dp_atomic.fitting_net.get_type_map(), data["type_map"])
71+
72+
73+
if __name__ == "__main__":
74+
unittest.main()

0 commit comments

Comments
 (0)