From c8de6a6cce857af4cb923fe0003adf5aebeaa6f1 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.2))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Mon, 16 Mar 2026 04:49:51 +0000 Subject: [PATCH 01/22] feat(show): add serialization-tree via DeepEval.serialize Authored by OpenClaw (model: gpt-5.2) --- deepmd/entrypoints/show.py | 18 ++++++++++++----- deepmd/infer/deep_eval.py | 41 ++++++++++++++++++++++++++++++++++++++ deepmd/main.py | 1 + 3 files changed, 55 insertions(+), 5 deletions(-) diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 7fd3e81467..0f6d023691 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -46,8 +46,7 @@ def show( f"where 'RANDOM' means using a randomly initialized fitting net." ) log.info( - "Detailed information: \n" - + OrderedDictTableWrapper(model_branch_dict).as_table() + "Detailed information: \n" + OrderedDictTableWrapper(model_branch_dict).as_table() ) if "type-map" in ATTRIBUTES: if model_is_multi_task: @@ -72,9 +71,7 @@ def show( model_branches = list(model_params["model_dict"].keys()) for branch in model_branches: fitting_net = model_params["model_dict"][branch]["fitting_net"] - log.info( - f"The fitting_net parameter of branch {branch} is {fitting_net}" - ) + log.info(f"The fitting_net parameter of branch {branch} is {fitting_net}") else: fitting_net = model_params["fitting_net"] log.info(f"The fitting_net parameter is {fitting_net}") @@ -136,3 +133,14 @@ def show( observed_types = model.get_observed_types() log.info(f"Number of observed types: {observed_types['type_num']} ") log.info(f"Observed types: {observed_types['observed_type']} ") + + if "serialization-tree" in ATTRIBUTES: + from deepmd.dpmodel.utils.serialization import ( + Node, + ) + + data = model.serialize() + if "model" not in data: + raise RuntimeError("Serialized model data does not contain key 'model'.") + root = Node.deserialize(data["model"]) + log.info("Model serialization tree:\n" + str(root)) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index d375a2ecd7..06a531c1e3 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -404,6 +404,7 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: + self.model_file = model_file self.deep_eval = DeepEvalBackend( model_file, self.output_def, @@ -420,6 +421,46 @@ def __init__( def output_def(self) -> ModelOutputDef: """Returns the output variable definitions.""" + def serialize(self) -> dict[str, Any]: + """Serialize the model file to a dictionary (backend-unified). + + This is a convenience wrapper around backend-specific serialization + hooks, intended for unified model inspection / display. + + Returns + ------- + dict + Serialized model data (must include key ``"model"``). + + Raises + ------ + NotImplementedError + If the detected backend does not support IO serialization. + """ + backend_cls = Backend.detect_backend_by_model(self.model_file) + + # internal alias backend: resolve to a verified local file first + if getattr(backend_cls, "name", "").lower() == "pretrained": + from deepmd.pretrained.deep_eval import ( + parse_pretrained_alias, + ) + from deepmd.pretrained.download import ( + resolve_model_path, + ) + + model_name = parse_pretrained_alias(self.model_file) + resolved = str(resolve_model_path(model_name)) + backend_cls = Backend.detect_backend_by_model(resolved) + return backend_cls().serialize_hook(resolved) + + if not (backend_cls.features & Backend.Feature.IO): + raise NotImplementedError( + f"Backend '{backend_cls.name}' does not support serialization." + ) + + return backend_cls().serialize_hook(self.model_file) + + def get_rcut(self) -> float: """Get the cutoff radius of this model.""" return self.deep_eval.get_rcut() diff --git a/deepmd/main.py b/deepmd/main.py index 3afcda8b4a..6ef2fcd0d2 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -942,6 +942,7 @@ def main_parser() -> argparse.ArgumentParser: "fitting-net", "size", "observed-type", + "serialization-tree", ], nargs="+", ) From b695aaa1ca4de6bb3b7f4b9377f1f190d4d0115b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 04:52:32 +0000 Subject: [PATCH 02/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/entrypoints/show.py | 7 +++++-- deepmd/infer/deep_eval.py | 1 - 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 0f6d023691..6ed2f0be02 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -46,7 +46,8 @@ def show( f"where 'RANDOM' means using a randomly initialized fitting net." ) log.info( - "Detailed information: \n" + OrderedDictTableWrapper(model_branch_dict).as_table() + "Detailed information: \n" + + OrderedDictTableWrapper(model_branch_dict).as_table() ) if "type-map" in ATTRIBUTES: if model_is_multi_task: @@ -71,7 +72,9 @@ def show( model_branches = list(model_params["model_dict"].keys()) for branch in model_branches: fitting_net = model_params["model_dict"][branch]["fitting_net"] - log.info(f"The fitting_net parameter of branch {branch} is {fitting_net}") + log.info( + f"The fitting_net parameter of branch {branch} is {fitting_net}" + ) else: fitting_net = model_params["fitting_net"] log.info(f"The fitting_net parameter is {fitting_net}") diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 06a531c1e3..b4d50da491 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -460,7 +460,6 @@ def serialize(self) -> dict[str, Any]: return backend_cls().serialize_hook(self.model_file) - def get_rcut(self) -> float: """Get the cutoff radius of this model.""" return self.deep_eval.get_rcut() From 1694360d30b973106ab306590004848bbd5a143b Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.2))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Mon, 16 Mar 2026 12:56:13 +0000 Subject: [PATCH 03/22] refactor(deepeval): implement backend serialize via model.serialize Switch DeepEval.serialize() to delegate to DeepEvalBackend.serialize(), and implement serialize() in each backend by calling the underlying model's serialize(). Also move Node import in dp show to module top-level. Authored by OpenClaw (model: gpt-5.2) --- deepmd/dpmodel/infer/deep_eval.py | 12 ++++++++ deepmd/entrypoints/show.py | 7 ++--- deepmd/infer/deep_eval.py | 49 ++++++++----------------------- deepmd/jax/infer/deep_eval.py | 16 ++++++++++ deepmd/pd/infer/deep_eval.py | 12 ++++++++ deepmd/pretrained/deep_eval.py | 3 ++ deepmd/pt/infer/deep_eval.py | 12 ++++++++ deepmd/pt_expt/infer/deep_eval.py | 7 +++++ deepmd/tf/infer/deep_eval.py | 34 +++++++++++++++++++++ 9 files changed, 111 insertions(+), 41 deletions(-) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index ac6963b435..27b583ef74 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -406,6 +406,18 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return json.loads(self.dp.get_model_def_script()) + def serialize(self) -> dict[str, Any]: + model = self.dp + data: dict[str, Any] = { + "backend": "DPModel", + "model": model.serialize(), + "model_def_script": self.get_model_def_script(), + "@variables": {}, + } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() + return data + def get_observed_types(self) -> dict: """Get observed types (elements) of the model during data statistics. diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 6ed2f0be02..e916a1d044 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -4,6 +4,9 @@ Any, ) +from deepmd.dpmodel.utils.serialization import ( + Node, +) from deepmd.infer.deep_eval import ( DeepEval, ) @@ -138,10 +141,6 @@ def show( log.info(f"Observed types: {observed_types['observed_type']} ") if "serialization-tree" in ATTRIBUTES: - from deepmd.dpmodel.utils.serialization import ( - Node, - ) - data = model.serialize() if "model" not in data: raise RuntimeError("Serialized model data does not contain key 'model'.") diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index b4d50da491..c8ebf1c472 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -361,6 +361,16 @@ def get_model(self) -> Any: The model module implemented by the deep learning framework. """ + @abstractmethod + def serialize(self) -> dict[str, Any]: + """Serialize the loaded model to a backend-unified dictionary. + + Returns + ------- + dict + Serialized model data. Must include key ``"model"``. + """ + class DeepEval(ABC): """High-level Deep Evaluator interface. @@ -422,43 +432,8 @@ def output_def(self) -> ModelOutputDef: """Returns the output variable definitions.""" def serialize(self) -> dict[str, Any]: - """Serialize the model file to a dictionary (backend-unified). - - This is a convenience wrapper around backend-specific serialization - hooks, intended for unified model inspection / display. - - Returns - ------- - dict - Serialized model data (must include key ``"model"``). - - Raises - ------ - NotImplementedError - If the detected backend does not support IO serialization. - """ - backend_cls = Backend.detect_backend_by_model(self.model_file) - - # internal alias backend: resolve to a verified local file first - if getattr(backend_cls, "name", "").lower() == "pretrained": - from deepmd.pretrained.deep_eval import ( - parse_pretrained_alias, - ) - from deepmd.pretrained.download import ( - resolve_model_path, - ) - - model_name = parse_pretrained_alias(self.model_file) - resolved = str(resolve_model_path(model_name)) - backend_cls = Backend.detect_backend_by_model(resolved) - return backend_cls().serialize_hook(resolved) - - if not (backend_cls.features & Backend.Feature.IO): - raise NotImplementedError( - f"Backend '{backend_cls.name}' does not support serialization." - ) - - return backend_cls().serialize_hook(self.model_file) + """Serialize the loaded model to a backend-unified dictionary.""" + return self.deep_eval.serialize() def get_rcut(self) -> float: """Get the cutoff radius of this model.""" diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 2e028225f7..d3dd6148b8 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -47,6 +47,9 @@ from deepmd.jax.common import ( to_jax_array, ) +from deepmd.jax.env import ( + jax, +) from deepmd.jax.model.hlo import ( HLO, ) @@ -187,6 +190,19 @@ def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" return 0 + def serialize(self) -> dict[str, Any]: + model = self.dp + data: dict[str, Any] = { + "backend": "JAX", + "jax_version": jax.__version__, + "model": model.serialize(), + "model_def_script": json.loads(model.get_model_def_script()), + "@variables": {}, + } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() + return data + def eval( self, coords: np.ndarray, diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 6c0ffed7ec..66ebe86621 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -730,6 +730,18 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return self.model_def_script + def serialize(self) -> dict[str, Any]: + model = self.dp.model["Default"] + data: dict[str, Any] = { + "backend": "Paddle", + "model": model.serialize(), + "model_def_script": self.get_model_def_script(), + "@variables": {}, + } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() + return data + def get_model_size(self) -> dict: """Get model parameter count. diff --git a/deepmd/pretrained/deep_eval.py b/deepmd/pretrained/deep_eval.py index 2dc671b0cc..aa15a50760 100644 --- a/deepmd/pretrained/deep_eval.py +++ b/deepmd/pretrained/deep_eval.py @@ -184,3 +184,6 @@ def get_ntypes_spin(self) -> int: def get_model(self) -> Any: return self._backend.get_model() + + def serialize(self) -> dict[str, Any]: + return self._backend.serialize() diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 11a877040d..720e44b9f5 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -702,6 +702,18 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return self.model_def_script + def serialize(self) -> dict[str, Any]: + model = self.dp.model["Default"] + data: dict[str, Any] = { + "backend": "PyTorch", + "model": model.serialize(), + "model_def_script": self.get_model_def_script(), + "@variables": {}, + } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() + return data + def get_model_size(self) -> dict: """Get model parameter count. diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index a6e1e1e540..992d43e8cd 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -665,6 +665,13 @@ def get_model_def_script(self) -> dict: """Get model definition script.""" return self.metadata + def serialize(self) -> dict[str, Any]: + from deepmd.pt_expt.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file(self.model_path) + def get_model(self) -> torch.nn.Module: """Get the exported model module. diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 0ec2f1c74e..afe8d6122f 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -110,6 +110,7 @@ def __init__( input_map=input_map, ) self.load_prefix = load_prefix + self.model_file = model_file # graph_compatable should be called after graph and prefix are set if not self._graph_compatable(): @@ -1121,6 +1122,38 @@ def get_model_def_script(self) -> dict: model_def_script = script.decode("utf-8") return json.loads(model_def_script)["model"] + def serialize(self) -> dict[str, Any]: + from deepmd.tf.model.model import ( + Model, + ) + from deepmd.tf.utils.graph import ( + load_graph_def, + ) + + graph, graph_def = load_graph_def(str(self.model_file)) + + model_def_script = self.get_model_def_script() + model = Model(**model_def_script) + # important! must be called before serialize + model.init_variables(graph=graph, graph_def=graph_def) + model_dict = model.serialize() + + data: dict[str, Any] = { + "backend": "TensorFlow", + "tf_version": tf.__version__, + "model": model_dict, + "model_def_script": model_def_script, + } + try: + t_min_nbor_dist = self._get_tensor("train_attr/min_nbor_dist:0") + except KeyError: + pass + else: + [min_nbor_dist] = run_sess(self.sess, [t_min_nbor_dist], feed_dict={}) + data.setdefault("@variables", {}) + data["@variables"]["min_nbor_dist"] = float(min_nbor_dist) + return data + def get_model(self) -> "tf.Graph": """Get the TensorFlow graph. @@ -1172,6 +1205,7 @@ def __init__( input_map=input_map, ) self.load_prefix = load_prefix + self.model_file = model_file # graph_compatable should be called after graph and prefix are set if not self._graph_compatable(): From 76753885d103fba7c0295042a5afdd51d1d40265 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Tue, 7 Apr 2026 17:46:55 +0000 Subject: [PATCH 04/22] fix(show): align backend serialize output contracts Route JAX DeepEval serialization through the existing file-based serializer so .hlo and .savedmodel models follow the supported path instead of calling unimplemented model-level serialize() methods. Also add the missing pt_version field to the PyTorch backend serializer and wrap pt_expt serialization in the backend-unified payload expected by dp show serialization-tree. Add a targeted pt_expt serialization contract test. Authored by OpenClaw (model: gpt-5.4) --- deepmd/jax/infer/deep_eval.py | 16 +++++----------- deepmd/pt/infer/deep_eval.py | 1 + deepmd/pt_expt/infer/deep_eval.py | 8 +++++++- source/tests/pt_expt/infer/test_deep_eval.py | 10 ++++++++++ 4 files changed, 23 insertions(+), 12 deletions(-) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index d3dd6148b8..aaf6e9900b 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -191,17 +191,11 @@ def get_ntypes_spin(self) -> int: return 0 def serialize(self) -> dict[str, Any]: - model = self.dp - data: dict[str, Any] = { - "backend": "JAX", - "jax_version": jax.__version__, - "model": model.serialize(), - "model_def_script": json.loads(model.get_model_def_script()), - "@variables": {}, - } - if model.get_min_nbor_dist() is not None: - data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() - return data + from deepmd.jax.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file(self.model_path) def eval( self, diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 720e44b9f5..aefc970637 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -706,6 +706,7 @@ def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] data: dict[str, Any] = { "backend": "PyTorch", + "pt_version": str(torch.__version__), "model": model.serialize(), "model_def_script": self.get_model_def_script(), "@variables": {}, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 992d43e8cd..2e9b8fefc9 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -670,7 +670,13 @@ def serialize(self) -> dict[str, Any]: serialize_from_file, ) - return serialize_from_file(self.model_path) + model_dict = serialize_from_file(self.model_path) + return { + "backend": "PyTorch Exportable", + "model": model_dict, + "model_def_script": self.get_model_def_script(), + "@variables": {}, + } def get_model(self) -> torch.nn.Module: """Get the exported model module. diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index ef38e1d36f..7a9273e3ab 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -109,6 +109,16 @@ def test_get_model_def_script(self) -> None: self.assertAlmostEqual(mds["rcut"], self.rcut) self.assertEqual(mds["sel"], list(self.sel)) + def test_serialize_contract(self) -> None: + data = self.dp.deep_eval.serialize() + self.assertEqual(data["backend"], "PyTorch Exportable") + self.assertIn("model", data) + self.assertIn("model_def_script", data) + self.assertIn("@variables", data) + self.assertIsInstance(data["@variables"], dict) + self.assertEqual(data["model_def_script"]["type_map"], self.type_map) + self.assertEqual(data["model"], serialize_from_file(self.tmpfile.name)) + def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" rng = np.random.default_rng(GLOBAL_SEED) From a8a80f4ff8f49d17b82e4af67548692e8b6e1bd4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 7 Apr 2026 17:48:01 +0000 Subject: [PATCH 05/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/jax/infer/deep_eval.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index aaf6e9900b..5c838de100 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -47,9 +47,6 @@ from deepmd.jax.common import ( to_jax_array, ) -from deepmd.jax.env import ( - jax, -) from deepmd.jax.model.hlo import ( HLO, ) From c18814cf9c92fda0617c07fbe5d7005444bdbe6f Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 02:39:06 +0000 Subject: [PATCH 06/22] refactor(show): decouple serialization tree from deep eval wrapper Route dp show serialization-tree through backend file serialization hooks and pass only the serialized model payload into Node.deserialize(). This keeps Node focused on model structure instead of backend-specific envelope fields. Add a focused unit test that locks the behavior at the entrypoint layer. Authored by OpenClaw (model: gpt-5.4) --- deepmd/entrypoints/show.py | 10 ++++-- ...test_entrypoint_show_serialization_tree.py | 33 +++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 source/tests/test_entrypoint_show_serialization_tree.py diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index e916a1d044..42d3ffbc92 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -4,6 +4,9 @@ Any, ) +from deepmd.backend.backend import ( + Backend, +) from deepmd.dpmodel.utils.serialization import ( Node, ) @@ -141,8 +144,11 @@ def show( log.info(f"Observed types: {observed_types['observed_type']} ") if "serialization-tree" in ATTRIBUTES: - data = model.serialize() + backend = Backend.detect_backend_by_model(INPUT)() + data = backend.serialize_hook(INPUT) if "model" not in data: - raise RuntimeError("Serialized model data does not contain key 'model'.") + raise RuntimeError( + "Serialized model data does not contain key 'model'." + ) root = Node.deserialize(data["model"]) log.info("Model serialization tree:\n" + str(root)) diff --git a/source/tests/test_entrypoint_show_serialization_tree.py b/source/tests/test_entrypoint_show_serialization_tree.py new file mode 100644 index 0000000000..299c60aec9 --- /dev/null +++ b/source/tests/test_entrypoint_show_serialization_tree.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from unittest.mock import patch + +from deepmd.entrypoints.show import show + + +class TestShowSerializationTree(unittest.TestCase): + def test_serialization_tree_uses_backend_hook_model_payload(self) -> None: + with ( + patch("deepmd.entrypoints.show.DeepEval") as mock_deep_eval, + patch("deepmd.entrypoints.show.Backend.detect_backend_by_model") as mock_detect, + patch("deepmd.entrypoints.show.Node.deserialize") as mock_deserialize, + patch("deepmd.entrypoints.show.log.info") as mock_log_info, + ): + model = mock_deep_eval.return_value + model.get_model_def_script.return_value = {"type_map": ["H", "O"]} + model.get_model_size.return_value = {} + + backend = mock_detect.return_value.return_value + backend.serialize_hook.return_value = { + "backend": "PyTorch Exportable", + "model": {"@class": "MockModel"}, + "model_def_script": {"type_map": ["H", "O"]}, + "@variables": {}, + } + mock_deserialize.return_value = "ROOT" + + show(INPUT="mock.pte", ATTRIBUTES=["serialization-tree"]) + + backend.serialize_hook.assert_called_once_with("mock.pte") + mock_deserialize.assert_called_once_with({"@class": "MockModel"}) + mock_log_info.assert_any_call("Model serialization tree:\nROOT") From 5678d784670b644aaf44d0af32e2ffcbde6ab6d2 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 02:49:47 +0000 Subject: [PATCH 07/22] refactor(deepeval): serialize model tree only Make DeepEval.serialize() return only the model serialization tree instead of a backend metadata envelope. This lets callers such as dp show serialization-tree consume model structure directly without caring about backend-specific data wrappers. Keep full backend data serialization in backend serialize hooks for IO and backend conversion paths. Authored by OpenClaw (model: gpt-5.4) --- deepmd/dpmodel/infer/deep_eval.py | 10 +--------- deepmd/entrypoints/show.py | 11 +---------- deepmd/infer/deep_eval.py | 6 +++--- deepmd/jax/infer/deep_eval.py | 5 ++++- deepmd/pd/infer/deep_eval.py | 10 +--------- deepmd/pt/infer/deep_eval.py | 11 +---------- deepmd/pt_expt/infer/deep_eval.py | 8 +------- deepmd/tf/infer/deep_eval.py | 18 +----------------- source/tests/pt_expt/infer/test_deep_eval.py | 12 ++++-------- .../test_entrypoint_show_serialization_tree.py | 14 +++----------- 10 files changed, 20 insertions(+), 85 deletions(-) diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 27b583ef74..d567ece0e1 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -408,15 +408,7 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp - data: dict[str, Any] = { - "backend": "DPModel", - "model": model.serialize(), - "model_def_script": self.get_model_def_script(), - "@variables": {}, - } - if model.get_min_nbor_dist() is not None: - data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() - return data + return model.serialize() def get_observed_types(self) -> dict: """Get observed types (elements) of the model during data statistics. diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 42d3ffbc92..2cf1b881bb 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -4,9 +4,6 @@ Any, ) -from deepmd.backend.backend import ( - Backend, -) from deepmd.dpmodel.utils.serialization import ( Node, ) @@ -144,11 +141,5 @@ def show( log.info(f"Observed types: {observed_types['observed_type']} ") if "serialization-tree" in ATTRIBUTES: - backend = Backend.detect_backend_by_model(INPUT)() - data = backend.serialize_hook(INPUT) - if "model" not in data: - raise RuntimeError( - "Serialized model data does not contain key 'model'." - ) - root = Node.deserialize(data["model"]) + root = Node.deserialize(model.serialize()) log.info("Model serialization tree:\n" + str(root)) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index c8ebf1c472..ae6720349e 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -363,12 +363,12 @@ def get_model(self) -> Any: @abstractmethod def serialize(self) -> dict[str, Any]: - """Serialize the loaded model to a backend-unified dictionary. + """Serialize the loaded model structure only. Returns ------- dict - Serialized model data. Must include key ``"model"``. + Serialized model tree that can be consumed by ``Node.deserialize``. """ @@ -432,7 +432,7 @@ def output_def(self) -> ModelOutputDef: """Returns the output variable definitions.""" def serialize(self) -> dict[str, Any]: - """Serialize the loaded model to a backend-unified dictionary.""" + """Serialize the loaded model structure only.""" return self.deep_eval.serialize() def get_rcut(self) -> float: diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 5c838de100..c2784c3b5f 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -192,7 +192,10 @@ def serialize(self) -> dict[str, Any]: serialize_from_file, ) - return serialize_from_file(self.model_path) + data = serialize_from_file(self.model_path) + if "model" not in data: + raise RuntimeError("Serialized model data does not contain key 'model'.") + return data["model"] def eval( self, diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 66ebe86621..8baef70c15 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -732,15 +732,7 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] - data: dict[str, Any] = { - "backend": "Paddle", - "model": model.serialize(), - "model_def_script": self.get_model_def_script(), - "@variables": {}, - } - if model.get_min_nbor_dist() is not None: - data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() - return data + return model.serialize() def get_model_size(self) -> dict: """Get model parameter count. diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index aefc970637..1e4c8a0929 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -704,16 +704,7 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] - data: dict[str, Any] = { - "backend": "PyTorch", - "pt_version": str(torch.__version__), - "model": model.serialize(), - "model_def_script": self.get_model_def_script(), - "@variables": {}, - } - if model.get_min_nbor_dist() is not None: - data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() - return data + return model.serialize() def get_model_size(self) -> dict: """Get model parameter count. diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 2e9b8fefc9..992d43e8cd 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -670,13 +670,7 @@ def serialize(self) -> dict[str, Any]: serialize_from_file, ) - model_dict = serialize_from_file(self.model_path) - return { - "backend": "PyTorch Exportable", - "model": model_dict, - "model_def_script": self.get_model_def_script(), - "@variables": {}, - } + return serialize_from_file(self.model_path) def get_model(self) -> torch.nn.Module: """Get the exported model module. diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index afe8d6122f..0d81ce588f 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1136,23 +1136,7 @@ def serialize(self) -> dict[str, Any]: model = Model(**model_def_script) # important! must be called before serialize model.init_variables(graph=graph, graph_def=graph_def) - model_dict = model.serialize() - - data: dict[str, Any] = { - "backend": "TensorFlow", - "tf_version": tf.__version__, - "model": model_dict, - "model_def_script": model_def_script, - } - try: - t_min_nbor_dist = self._get_tensor("train_attr/min_nbor_dist:0") - except KeyError: - pass - else: - [min_nbor_dist] = run_sess(self.sess, [t_min_nbor_dist], feed_dict={}) - data.setdefault("@variables", {}) - data["@variables"]["min_nbor_dist"] = float(min_nbor_dist) - return data + return model.serialize() def get_model(self) -> "tf.Graph": """Get the TensorFlow graph. diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 7a9273e3ab..26d0741f7b 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -109,15 +109,11 @@ def test_get_model_def_script(self) -> None: self.assertAlmostEqual(mds["rcut"], self.rcut) self.assertEqual(mds["sel"], list(self.sel)) - def test_serialize_contract(self) -> None: + def test_serialize_returns_model_tree(self) -> None: data = self.dp.deep_eval.serialize() - self.assertEqual(data["backend"], "PyTorch Exportable") - self.assertIn("model", data) - self.assertIn("model_def_script", data) - self.assertIn("@variables", data) - self.assertIsInstance(data["@variables"], dict) - self.assertEqual(data["model_def_script"]["type_map"], self.type_map) - self.assertEqual(data["model"], serialize_from_file(self.tmpfile.name)) + self.assertEqual(data["@class"], self.model.serialize()["@class"]) + self.assertEqual(data["type"], self.model.serialize()["type"]) + self.assertEqual(data, serialize_from_file(self.tmpfile.name)) def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" diff --git a/source/tests/test_entrypoint_show_serialization_tree.py b/source/tests/test_entrypoint_show_serialization_tree.py index 299c60aec9..f7a90fb887 100644 --- a/source/tests/test_entrypoint_show_serialization_tree.py +++ b/source/tests/test_entrypoint_show_serialization_tree.py @@ -6,28 +6,20 @@ class TestShowSerializationTree(unittest.TestCase): - def test_serialization_tree_uses_backend_hook_model_payload(self) -> None: + def test_serialization_tree_uses_deep_eval_model_payload(self) -> None: with ( patch("deepmd.entrypoints.show.DeepEval") as mock_deep_eval, - patch("deepmd.entrypoints.show.Backend.detect_backend_by_model") as mock_detect, patch("deepmd.entrypoints.show.Node.deserialize") as mock_deserialize, patch("deepmd.entrypoints.show.log.info") as mock_log_info, ): model = mock_deep_eval.return_value model.get_model_def_script.return_value = {"type_map": ["H", "O"]} model.get_model_size.return_value = {} - - backend = mock_detect.return_value.return_value - backend.serialize_hook.return_value = { - "backend": "PyTorch Exportable", - "model": {"@class": "MockModel"}, - "model_def_script": {"type_map": ["H", "O"]}, - "@variables": {}, - } + model.serialize.return_value = {"@class": "MockModel"} mock_deserialize.return_value = "ROOT" show(INPUT="mock.pte", ATTRIBUTES=["serialization-tree"]) - backend.serialize_hook.assert_called_once_with("mock.pte") + model.serialize.assert_called_once_with() mock_deserialize.assert_called_once_with({"@class": "MockModel"}) mock_log_info.assert_any_call("Model serialization tree:\nROOT") From 8b48f46ac7d2e1a6d89b352fbace400af12fad81 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 03:01:09 +0000 Subject: [PATCH 08/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/test_entrypoint_show_serialization_tree.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/source/tests/test_entrypoint_show_serialization_tree.py b/source/tests/test_entrypoint_show_serialization_tree.py index f7a90fb887..9f01651115 100644 --- a/source/tests/test_entrypoint_show_serialization_tree.py +++ b/source/tests/test_entrypoint_show_serialization_tree.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest -from unittest.mock import patch +from unittest.mock import ( + patch, +) -from deepmd.entrypoints.show import show +from deepmd.entrypoints.show import ( + show, +) class TestShowSerializationTree(unittest.TestCase): From 5e97073b229820fe535a0556f99594fbca53a132 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 03:18:20 +0000 Subject: [PATCH 09/22] test(io): cover deep eval serialization in consistent io Assert in the consistent IO deep-eval test that DeepEval.serialize() returns the same model tree as the backend serialize hook payload. This keeps the new serialize semantics covered in the shared cross-backend IO test. Authored by OpenClaw (model: gpt-5.4) --- source/tests/consistent/io/test_io.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 982d56d8fa..8e5df1ba58 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -156,11 +156,12 @@ def test_deep_eval(self) -> None: if not backend.is_available(): continue reference_data = copy.deepcopy(self.data) - self.save_data_to_model( - prefix + backend.suffixes[suffix_idx], reference_data - ) - deep_eval = DeepEval(prefix + backend.suffixes[suffix_idx]) + model_file = prefix + backend.suffixes[suffix_idx] + self.save_data_to_model(model_file, reference_data) + deep_eval = DeepEval(model_file) self.assertIsInstance(deep_eval.get_model_def_script(), dict) + serialized_data = self.get_data_from_model(model_file) + self.assertEqual(deep_eval.serialize(), serialized_data["model"]) if deep_eval.get_dim_fparam() > 0: fparam = np.ones((nframes, deep_eval.get_dim_fparam())) else: From 8bf317624ef629f3f5dde041d3fa53491f92a60b Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 03:33:12 +0000 Subject: [PATCH 10/22] fix(review): handle paddle static serialize and io assert Handle Paddle static models when serializing through DeepEval by falling back to self.dp when it is not a ModelWrapper. Also use numpy-aware equality in the consistent IO serialize check. Authored by OpenClaw (model: gpt-5.4) --- deepmd/pd/infer/deep_eval.py | 2 +- source/tests/consistent/io/test_io.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 8baef70c15..65af5e0884 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -731,7 +731,7 @@ def get_model_def_script(self) -> dict: return self.model_def_script def serialize(self) -> dict[str, Any]: - model = self.dp.model["Default"] + model = self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp return model.serialize() def get_model_size(self) -> dict: diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 8e5df1ba58..c492e52510 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -161,7 +161,7 @@ def test_deep_eval(self) -> None: deep_eval = DeepEval(model_file) self.assertIsInstance(deep_eval.get_model_def_script(), dict) serialized_data = self.get_data_from_model(model_file) - self.assertEqual(deep_eval.serialize(), serialized_data["model"]) + np.testing.assert_equal(deep_eval.serialize(), serialized_data["model"]) if deep_eval.get_dim_fparam() > 0: fparam = np.ones((nframes, deep_eval.get_dim_fparam())) else: From 6cf9dbf577886dcc55ba819fb19c4fd7cb71171d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 03:33:59 +0000 Subject: [PATCH 11/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pd/infer/deep_eval.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 65af5e0884..561faca562 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -731,7 +731,9 @@ def get_model_def_script(self) -> dict: return self.model_def_script def serialize(self) -> dict[str, Any]: - model = self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp + model = ( + self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp + ) return model.serialize() def get_model_size(self) -> dict: From 53f6f12373e029a8ab8410068a87b35144c38d95 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 06:22:32 +0000 Subject: [PATCH 12/22] fix(ci): handle savedmodel and torchscript serialize --- deepmd/jax/utils/serialization.py | 10 +++++++++- deepmd/pt/infer/deep_eval.py | 9 ++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 14386d9f3d..ba6ad31ea0 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -202,5 +202,13 @@ def convert_str_to_int_key(item: dict) -> None: data.pop("constants") data["@variables"].pop("stablehlo") return data + elif model_file.endswith(".savedmodel"): + from deepmd.tf.utils.serialization import ( + serialize_from_file as serialize_savedmodel, + ) + + return serialize_savedmodel(model_file) else: - raise ValueError("JAX backend only supports converting .jax directory") + raise ValueError( + "JAX backend only supports converting .jax directory, .hlo, and .savedmodel" + ) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 1e4c8a0929..5418a67c03 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -704,7 +704,14 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] - return model.serialize() + if hasattr(model, "serialize"): + return model.serialize() + + from deepmd.pt.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file(self.model_path)["model"] def get_model_size(self) -> dict: """Get model parameter count. From a2431e8fcdbfba89aebf795ed124647e720df56d Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 09:02:51 +0000 Subject: [PATCH 13/22] fix(ci): guard pt serialize fallback shape --- deepmd/pt/infer/deep_eval.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 5418a67c03..c13659968a 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -711,7 +711,8 @@ def serialize(self) -> dict[str, Any]: serialize_from_file, ) - return serialize_from_file(self.model_path)["model"] + data = serialize_from_file(self.model_path) + return data["model"] if isinstance(data, dict) and "model" in data else data def get_model_size(self) -> dict: """Get model parameter count. From 39aab42cd4a3eb7071d6cbf4d5699c9b11aed3f6 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 11:02:19 +0000 Subject: [PATCH 14/22] fix(test): pt_expt serialize fallback returns model tree with @class key --- deepmd/pt/infer/deep_eval.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index c13659968a..e6855c1236 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -707,11 +707,26 @@ def serialize(self) -> dict[str, Any]: if hasattr(model, "serialize"): return model.serialize() + # Try pt_expt serialization first (for .pte files) + try: + from deepmd.pt_expt.utils.serialization import ( + serialize_from_file as serialize_from_pte, + ) + data = serialize_from_pte(self.model_path) + # pt_expt serialize_from_file returns the model tree directly + return data + except ImportError: + pass + except Exception: + # If pt_expt serialization fails for other reasons, fall back + pass + + # Fallback to generic pt serialization (for .pth/.pt files) from deepmd.pt.utils.serialization import ( serialize_from_file, ) - data = serialize_from_file(self.model_path) + # Generic pt serialize_from_file returns wrapped dict with "model" key return data["model"] if isinstance(data, dict) and "model" in data else data def get_model_size(self) -> dict: From 67f4d1e77bedde1c9b145e429cf4924223198916 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 11:03:19 +0000 Subject: [PATCH 15/22] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/infer/deep_eval.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index e6855c1236..4ddf10ac3f 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -712,6 +712,7 @@ def serialize(self) -> dict[str, Any]: from deepmd.pt_expt.utils.serialization import ( serialize_from_file as serialize_from_pte, ) + data = serialize_from_pte(self.model_path) # pt_expt serialize_from_file returns the model tree directly return data @@ -725,6 +726,7 @@ def serialize(self) -> dict[str, Any]: from deepmd.pt.utils.serialization import ( serialize_from_file, ) + data = serialize_from_file(self.model_path) # Generic pt serialize_from_file returns wrapped dict with "model" key return data["model"] if isinstance(data, dict) and "model" in data else data From 024a058a222795c0fa1553f44c097d64095cb7fe Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: gpt-5.4))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 13:24:19 +0000 Subject: [PATCH 16/22] fix(ci): revert pt deep eval mixed pte serialization path --- deepmd/pt/infer/deep_eval.py | 27 +-------------------------- 1 file changed, 1 insertion(+), 26 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 4ddf10ac3f..1e4c8a0929 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -704,32 +704,7 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] - if hasattr(model, "serialize"): - return model.serialize() - - # Try pt_expt serialization first (for .pte files) - try: - from deepmd.pt_expt.utils.serialization import ( - serialize_from_file as serialize_from_pte, - ) - - data = serialize_from_pte(self.model_path) - # pt_expt serialize_from_file returns the model tree directly - return data - except ImportError: - pass - except Exception: - # If pt_expt serialization fails for other reasons, fall back - pass - - # Fallback to generic pt serialization (for .pth/.pt files) - from deepmd.pt.utils.serialization import ( - serialize_from_file, - ) - - data = serialize_from_file(self.model_path) - # Generic pt serialize_from_file returns wrapped dict with "model" key - return data["model"] if isinstance(data, dict) and "model" in data else data + return model.serialize() def get_model_size(self) -> dict: """Get model parameter count. From b951333bb5d9ebe6eee78c41c15282e4a384a74e Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: kimi-k2.5))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Wed, 8 Apr 2026 15:45:29 +0000 Subject: [PATCH 17/22] fix(ci): restore backend-specific serialize fallbacks --- deepmd/pt/infer/deep_eval.py | 9 ++++++++- deepmd/pt_expt/infer/deep_eval.py | 3 ++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 1e4c8a0929..5418a67c03 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -704,7 +704,14 @@ def get_model_def_script(self) -> dict: def serialize(self) -> dict[str, Any]: model = self.dp.model["Default"] - return model.serialize() + if hasattr(model, "serialize"): + return model.serialize() + + from deepmd.pt.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file(self.model_path)["model"] def get_model_size(self) -> dict: """Get model parameter count. diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 992d43e8cd..8a3c6a7fe0 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -670,7 +670,8 @@ def serialize(self) -> dict[str, Any]: serialize_from_file, ) - return serialize_from_file(self.model_path) + data = serialize_from_file(self.model_path) + return data["model"] if isinstance(data, dict) and "model" in data else data def get_model(self) -> torch.nn.Module: """Get the exported model module. From f65bb5f1f9689c04148becc1201a5a371d80a3c1 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 21 Jun 2026 15:10:19 +0800 Subject: [PATCH 18/22] fix(jax): avoid lossless savedmodel serialization --- deepmd/jax/infer/deep_eval.py | 7 +++++++ deepmd/jax/utils/serialization.py | 10 +++++----- source/tests/consistent/io/test_io.py | 6 ++++-- source/tests/pt_expt/infer/test_deep_eval.py | 2 +- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index c2784c3b5f..44fd6d4c0d 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -188,6 +188,13 @@ def get_ntypes_spin(self) -> int: return 0 def serialize(self) -> dict[str, Any]: + if str(self.model_path).endswith(".savedmodel"): + from deepmd.jax.model.model import ( + get_model, + ) + + return get_model(self.get_model_def_script()).serialize() + from deepmd.jax.utils.serialization import ( serialize_from_file, ) diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index ba6ad31ea0..eded22b692 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -203,12 +203,12 @@ def convert_str_to_int_key(item: dict) -> None: data["@variables"].pop("stablehlo") return data elif model_file.endswith(".savedmodel"): - from deepmd.tf.utils.serialization import ( - serialize_from_file as serialize_savedmodel, + raise ValueError( + "JAX SavedModel does not support lossless file serialization. " + "Use DeepEval.serialize() for a structure-only model tree." ) - - return serialize_savedmodel(model_file) else: raise ValueError( - "JAX backend only supports converting .jax directory, .hlo, and .savedmodel" + "JAX backend only supports lossless file serialization for .jax " + "directory and .hlo." ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index c492e52510..574b3c472f 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -160,8 +160,10 @@ def test_deep_eval(self) -> None: self.save_data_to_model(model_file, reference_data) deep_eval = DeepEval(model_file) self.assertIsInstance(deep_eval.get_model_def_script(), dict) - serialized_data = self.get_data_from_model(model_file) - np.testing.assert_equal(deep_eval.serialize(), serialized_data["model"]) + if not model_file.endswith(".savedmodel"): + # JAX SavedModel stores an executable graph, not a lossless model dict. + serialized_data = self.get_data_from_model(model_file) + np.testing.assert_equal(deep_eval.serialize(), serialized_data["model"]) if deep_eval.get_dim_fparam() > 0: fparam = np.ones((nframes, deep_eval.get_dim_fparam())) else: diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 26d0741f7b..f5bac1d813 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -113,7 +113,7 @@ def test_serialize_returns_model_tree(self) -> None: data = self.dp.deep_eval.serialize() self.assertEqual(data["@class"], self.model.serialize()["@class"]) self.assertEqual(data["type"], self.model.serialize()["type"]) - self.assertEqual(data, serialize_from_file(self.tmpfile.name)) + self.assertEqual(data, serialize_from_file(self.tmpfile.name)["model"]) def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" From ebda14db231b626a455c93e505e49367c4c0916b Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Sun, 21 Jun 2026 09:06:42 +0000 Subject: [PATCH 19/22] test(pt-expt): compare serialized trees with numpy-aware assert Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5) --- source/tests/pt_expt/infer/test_deep_eval.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index f5bac1d813..3adb48833f 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -113,7 +113,10 @@ def test_serialize_returns_model_tree(self) -> None: data = self.dp.deep_eval.serialize() self.assertEqual(data["@class"], self.model.serialize()["@class"]) self.assertEqual(data["type"], self.model.serialize()["type"]) - self.assertEqual(data, serialize_from_file(self.tmpfile.name)["model"]) + # The serialized model tree contains NumPy array leaves, so unittest's + # dict equality would try to coerce elementwise array comparisons to a + # single bool and fail with an ambiguous truth-value error. + np.testing.assert_equal(data, serialize_from_file(self.tmpfile.name)["model"]) def test_eval_consistency(self) -> None: """Test that DeepPot.eval gives same results as direct model forward.""" From 1140851b13da77c9d2c4baea627e19a1fa548283 Mon Sep 17 00:00:00 2001 From: "njzjz-bot (driven by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5))[bot]" <48687836+njzjz-bot@users.noreply.github.com> Date: Mon, 22 Jun 2026 19:37:47 +0000 Subject: [PATCH 20/22] fix: address serialization tree review comments Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5) --- deepmd/entrypoints/show.py | 17 +- deepmd/infer/deep_eval.py | 23 ++- deepmd/jax/infer/deep_eval.py | 8 + deepmd/pd/infer/deep_eval.py | 9 +- source/tests/consistent/io/test_io.py | 2 + source/tests/test_deep_eval_serialize_api.py | 147 ++++++++++++++++++ ...test_entrypoint_show_serialization_tree.py | 48 +++++- 7 files changed, 243 insertions(+), 11 deletions(-) create mode 100644 source/tests/test_deep_eval_serialize_api.py diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 29e2a2384b..002fb159c4 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -4,9 +4,6 @@ Any, ) -from deepmd.dpmodel.utils.serialization import ( - Node, -) from deepmd.infer.deep_eval import ( DeepEval, ) @@ -154,5 +151,15 @@ def show( log.info(f"Observed types: {observed_types['observed_type']} ") if "serialization-tree" in ATTRIBUTES: - root = Node.deserialize(model.serialize()) - log.info("Model serialization tree:\n" + str(root)) + from deepmd.dpmodel.utils.serialization import ( + Node, + ) + + if model_is_multi_task: + for branch in model_params["model_dict"]: + branch_model = DeepEval(INPUT, head=branch) + root = Node.deserialize(branch_model.serialize()) + log.info("Model serialization tree of branch %s:\n%s", branch, root) + else: + root = Node.deserialize(model.serialize()) + log.info("Model serialization tree:\n%s", root) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 521c7fed9c..a08ec885f3 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -380,15 +380,26 @@ def get_model(self) -> Any: The model module implemented by the deep learning framework. """ - @abstractmethod def serialize(self) -> dict[str, Any]: - """Serialize the loaded model structure only. + """Serialize the loaded model as a model tree. + + Most in-tree backends return the lossless, weight-bearing ``model`` + subtree from the serialized file payload. Backends that cannot recover a + lossless tree may override this method to document and implement their + narrower behavior. Returns ------- dict Serialized model tree that can be consumed by ``Node.deserialize``. """ + model = self.get_model() + if hasattr(model, "serialize"): + return model.serialize() + raise NotImplementedError( + f"{type(self).__name__} does not implement serialize(), and its " + "model object has no serialize() method." + ) class DeepEval(ABC): @@ -451,7 +462,13 @@ def output_def(self) -> ModelOutputDef: """Returns the output variable definitions.""" def serialize(self) -> dict[str, Any]: - """Serialize the loaded model structure only.""" + """Serialize the loaded model as a model tree. + + Most backends return the lossless, weight-bearing ``model`` subtree from + the serialized file payload. JAX ``.savedmodel`` inputs are the known + exception: they are reconstructed from the model definition script and + therefore do not preserve trained weights. + """ return self.deep_eval.serialize() def get_rcut(self) -> float: diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 44fd6d4c0d..09b5783ba9 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -188,6 +188,14 @@ def get_ntypes_spin(self) -> int: return 0 def serialize(self) -> dict[str, Any]: + """Serialize the loaded model as a model tree. + + JAX-native ``.jax``/``.hlo`` inputs return the lossless, weight-bearing + ``model`` subtree from the file payload. TensorFlow-wrapped + ``.savedmodel`` inputs cannot be converted back losslessly; for that + format this method reconstructs the model tree from the definition + script, so trained weights are not preserved. + """ if str(self.model_path).endswith(".savedmodel"): from deepmd.jax.model.model import ( get_model, diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 2072b06a79..c8bb113495 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -744,7 +744,14 @@ def serialize(self) -> dict[str, Any]: model = ( self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp ) - return model.serialize() + if hasattr(model, "serialize"): + return model.serialize() + + from deepmd.pd.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file(self.model_path)["model"] def get_model_size(self) -> dict: """Get model parameter count. diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 574b3c472f..1b98955b9d 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -78,6 +78,7 @@ def test_data_equal(self) -> None: for backend_name, suffix_idx in ( ("tensorflow", 0) if not DP_TEST_TF2_ONLY else ("jax", 0), ("pytorch", 0), + ("paddle", 0), ("dpmodel", 0), ): with self.subTest(backend_name=backend_name): @@ -147,6 +148,7 @@ def test_deep_eval(self) -> None: # unfortunately, jax2tf cannot work with tf v1 behaviors ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), ("pytorch", 0), + ("paddle", 0), ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), ): diff --git a/source/tests/test_deep_eval_serialize_api.py b/source/tests/test_deep_eval_serialize_api.py new file mode 100644 index 0000000000..819c6f7486 --- /dev/null +++ b/source/tests/test_deep_eval_serialize_api.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from unittest.mock import ( + Mock, + patch, +) + +from deepmd.infer.deep_eval import ( + DeepEvalBackend, +) + + +class _DefaultSerializeBackend(DeepEvalBackend): + def __init__(self, model: object) -> None: + self._model = model + + def eval(self, *args: object, **kwargs: object) -> dict: + return {} + + def get_rcut(self) -> float: + return 0.0 + + def get_ntypes(self) -> int: + return 0 + + def get_type_map(self) -> list[str]: + return [] + + def get_dim_fparam(self) -> int: + return 0 + + def get_dim_aparam(self) -> int: + return 0 + + @property + def model_type(self) -> type: + return object + + def get_sel_type(self) -> list[int]: + return [] + + def get_ntypes_spin(self) -> int: + return 0 + + def get_model(self) -> object: + return self._model + + +class TestDeepEvalBackendSerialize(unittest.TestCase): + def test_default_serialize_delegates_to_model_when_available(self) -> None: + model = Mock() + model.serialize.return_value = {"@class": "Model"} + backend = _DefaultSerializeBackend(model) + + self.assertEqual(backend.serialize(), {"@class": "Model"}) + model.serialize.assert_called_once_with() + + def test_default_serialize_has_clear_error_without_model_method(self) -> None: + backend = _DefaultSerializeBackend(object()) + + with self.assertRaisesRegex( + NotImplementedError, "does not implement serialize" + ): + backend.serialize() + + +def _load_deep_eval_backend(module_name: str, backend_name: str): + try: + module = __import__(module_name, fromlist=["DeepEval"]) + except ImportError as exc: + raise unittest.SkipTest( + f"{backend_name} backend is not importable: {exc}" + ) from exc + return module.DeepEval + + +class TestPaddleDeepEvalSerialize(unittest.TestCase): + def test_jit_model_falls_back_to_file_serializer(self) -> None: + PaddleDeepEvalBackend = _load_deep_eval_backend( + "deepmd.pd.infer.deep_eval", "Paddle" + ) + backend = PaddleDeepEvalBackend.__new__(PaddleDeepEvalBackend) + backend.model_path = "frozen_model.json" + backend.dp = object() + + with patch("deepmd.pd.utils.serialization.serialize_from_file") as serialize: + serialize.return_value = {"model": {"@class": "RecoveredModel"}} + + self.assertEqual(backend.serialize(), {"@class": "RecoveredModel"}) + + serialize.assert_called_once_with("frozen_model.json") + + +class TestPyTorchDeepEvalSerialize(unittest.TestCase): + def test_jit_model_falls_back_to_file_serializer(self) -> None: + PyTorchDeepEvalBackend = _load_deep_eval_backend( + "deepmd.pt.infer.deep_eval", "PyTorch" + ) + backend = PyTorchDeepEvalBackend.__new__(PyTorchDeepEvalBackend) + backend.model_path = "frozen_model.pth" + backend.dp = Mock() + backend.dp.model = {"Default": object()} + + with patch("deepmd.pt.utils.serialization.serialize_from_file") as serialize: + serialize.return_value = {"model": {"@class": "RecoveredModel"}} + + self.assertEqual(backend.serialize(), {"@class": "RecoveredModel"}) + + serialize.assert_called_once_with("frozen_model.pth") + + +class TestPyTorchExportableDeepEvalSerialize(unittest.TestCase): + def test_raw_model_payload_fallback_is_preserved(self) -> None: + PyTorchExportableDeepEvalBackend = _load_deep_eval_backend( + "deepmd.pt_expt.infer.deep_eval", "PyTorch exportable" + ) + backend = PyTorchExportableDeepEvalBackend.__new__( + PyTorchExportableDeepEvalBackend + ) + backend.model_path = "frozen_model.pt" + + with patch( + "deepmd.pt_expt.utils.serialization.serialize_from_file" + ) as serialize: + serialize.return_value = {"@class": "RawExportedModel"} + + self.assertEqual(backend.serialize(), {"@class": "RawExportedModel"}) + + serialize.assert_called_once_with("frozen_model.pt") + + +class TestJAXDeepEvalSerialize(unittest.TestCase): + def test_savedmodel_reconstructs_tree_from_model_def_script(self) -> None: + JAXDeepEvalBackend = _load_deep_eval_backend( + "deepmd.jax.infer.deep_eval", "JAX" + ) + backend = JAXDeepEvalBackend.__new__(JAXDeepEvalBackend) + backend.model_path = "frozen_model.savedmodel" + backend.get_model_def_script = Mock(return_value={"type_map": ["O", "H"]}) + + model = Mock() + model.serialize.return_value = {"@class": "SavedModelTree"} + with patch("deepmd.jax.model.model.get_model", return_value=model) as get_model: + self.assertEqual(backend.serialize(), {"@class": "SavedModelTree"}) + + get_model.assert_called_once_with({"type_map": ["O", "H"]}) + model.serialize.assert_called_once_with() diff --git a/source/tests/test_entrypoint_show_serialization_tree.py b/source/tests/test_entrypoint_show_serialization_tree.py index 9f01651115..ddc30921f1 100644 --- a/source/tests/test_entrypoint_show_serialization_tree.py +++ b/source/tests/test_entrypoint_show_serialization_tree.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest from unittest.mock import ( + Mock, + call, patch, ) @@ -13,7 +15,9 @@ class TestShowSerializationTree(unittest.TestCase): def test_serialization_tree_uses_deep_eval_model_payload(self) -> None: with ( patch("deepmd.entrypoints.show.DeepEval") as mock_deep_eval, - patch("deepmd.entrypoints.show.Node.deserialize") as mock_deserialize, + patch( + "deepmd.dpmodel.utils.serialization.Node.deserialize" + ) as mock_deserialize, patch("deepmd.entrypoints.show.log.info") as mock_log_info, ): model = mock_deep_eval.return_value @@ -26,4 +30,44 @@ def test_serialization_tree_uses_deep_eval_model_payload(self) -> None: model.serialize.assert_called_once_with() mock_deserialize.assert_called_once_with({"@class": "MockModel"}) - mock_log_info.assert_any_call("Model serialization tree:\nROOT") + mock_log_info.assert_any_call("Model serialization tree:\n%s", "ROOT") + + def test_serialization_tree_iterates_multitask_branches(self) -> None: + with ( + patch("deepmd.entrypoints.show.DeepEval") as mock_deep_eval, + patch( + "deepmd.dpmodel.utils.serialization.Node.deserialize" + ) as mock_deserialize, + patch("deepmd.entrypoints.show.log.info") as mock_log_info, + ): + initial_model = mock_deep_eval.return_value + branch_a_model = mock_deep_eval.return_value + branch_b_model = Mock() + mock_deep_eval.side_effect = [initial_model, branch_a_model, branch_b_model] + + initial_model.get_model_def_script.return_value = { + "model_dict": {"branch_a": {}, "branch_b": {}} + } + initial_model.get_model_size.return_value = {} + branch_a_model.serialize.return_value = {"@class": "BranchA"} + branch_b_model.serialize.return_value = {"@class": "BranchB"} + mock_deserialize.side_effect = ["ROOT_A", "ROOT_B"] + + show(INPUT="mock-multitask.pte", ATTRIBUTES=["serialization-tree"]) + + mock_deep_eval.assert_has_calls( + [ + call("mock-multitask.pte", head=0), + call("mock-multitask.pte", head="branch_a"), + call("mock-multitask.pte", head="branch_b"), + ] + ) + mock_deserialize.assert_has_calls( + [call({"@class": "BranchA"}), call({"@class": "BranchB"})] + ) + mock_log_info.assert_any_call( + "Model serialization tree of branch %s:\n%s", "branch_a", "ROOT_A" + ) + mock_log_info.assert_any_call( + "Model serialization tree of branch %s:\n%s", "branch_b", "ROOT_B" + ) From 5e34f15e4c4624983e3357888103ae23f65c92be Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 28 Jun 2026 22:58:21 +0800 Subject: [PATCH 21/22] test(show): fix serialization tree CI failures --- source/tests/consistent/io/test_io.py | 2 -- source/tests/test_deep_eval_serialize_api.py | 10 ++++------ .../tests/test_entrypoint_show_serialization_tree.py | 9 +++++---- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 1b98955b9d..574b3c472f 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -78,7 +78,6 @@ def test_data_equal(self) -> None: for backend_name, suffix_idx in ( ("tensorflow", 0) if not DP_TEST_TF2_ONLY else ("jax", 0), ("pytorch", 0), - ("paddle", 0), ("dpmodel", 0), ): with self.subTest(backend_name=backend_name): @@ -148,7 +147,6 @@ def test_deep_eval(self) -> None: # unfortunately, jax2tf cannot work with tf v1 behaviors ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), ("pytorch", 0), - ("paddle", 0), ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), ): diff --git a/source/tests/test_deep_eval_serialize_api.py b/source/tests/test_deep_eval_serialize_api.py index 819c6f7486..6959a1bcb3 100644 --- a/source/tests/test_deep_eval_serialize_api.py +++ b/source/tests/test_deep_eval_serialize_api.py @@ -79,7 +79,7 @@ def test_jit_model_falls_back_to_file_serializer(self) -> None: PaddleDeepEvalBackend = _load_deep_eval_backend( "deepmd.pd.infer.deep_eval", "Paddle" ) - backend = PaddleDeepEvalBackend.__new__(PaddleDeepEvalBackend) + backend = object.__new__(PaddleDeepEvalBackend) backend.model_path = "frozen_model.json" backend.dp = object() @@ -96,7 +96,7 @@ def test_jit_model_falls_back_to_file_serializer(self) -> None: PyTorchDeepEvalBackend = _load_deep_eval_backend( "deepmd.pt.infer.deep_eval", "PyTorch" ) - backend = PyTorchDeepEvalBackend.__new__(PyTorchDeepEvalBackend) + backend = object.__new__(PyTorchDeepEvalBackend) backend.model_path = "frozen_model.pth" backend.dp = Mock() backend.dp.model = {"Default": object()} @@ -114,9 +114,7 @@ def test_raw_model_payload_fallback_is_preserved(self) -> None: PyTorchExportableDeepEvalBackend = _load_deep_eval_backend( "deepmd.pt_expt.infer.deep_eval", "PyTorch exportable" ) - backend = PyTorchExportableDeepEvalBackend.__new__( - PyTorchExportableDeepEvalBackend - ) + backend = object.__new__(PyTorchExportableDeepEvalBackend) backend.model_path = "frozen_model.pt" with patch( @@ -134,7 +132,7 @@ def test_savedmodel_reconstructs_tree_from_model_def_script(self) -> None: JAXDeepEvalBackend = _load_deep_eval_backend( "deepmd.jax.infer.deep_eval", "JAX" ) - backend = JAXDeepEvalBackend.__new__(JAXDeepEvalBackend) + backend = object.__new__(JAXDeepEvalBackend) backend.model_path = "frozen_model.savedmodel" backend.get_model_def_script = Mock(return_value={"type_map": ["O", "H"]}) diff --git a/source/tests/test_entrypoint_show_serialization_tree.py b/source/tests/test_entrypoint_show_serialization_tree.py index ddc30921f1..839a6b48a7 100644 --- a/source/tests/test_entrypoint_show_serialization_tree.py +++ b/source/tests/test_entrypoint_show_serialization_tree.py @@ -40,8 +40,8 @@ def test_serialization_tree_iterates_multitask_branches(self) -> None: ) as mock_deserialize, patch("deepmd.entrypoints.show.log.info") as mock_log_info, ): - initial_model = mock_deep_eval.return_value - branch_a_model = mock_deep_eval.return_value + initial_model = Mock() + branch_a_model = Mock() branch_b_model = Mock() mock_deep_eval.side_effect = [initial_model, branch_a_model, branch_b_model] @@ -55,12 +55,13 @@ def test_serialization_tree_iterates_multitask_branches(self) -> None: show(INPUT="mock-multitask.pte", ATTRIBUTES=["serialization-tree"]) - mock_deep_eval.assert_has_calls( + self.assertEqual( + mock_deep_eval.call_args_list, [ call("mock-multitask.pte", head=0), call("mock-multitask.pte", head="branch_a"), call("mock-multitask.pte", head="branch_b"), - ] + ], ) mock_deserialize.assert_has_calls( [call({"@class": "BranchA"}), call({"@class": "BranchB"})] From 13739768e0bdf7091949ae1cca3afe1a8357f009 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 1 Jul 2026 14:26:42 +0800 Subject: [PATCH 22/22] test(io): skip TF2 SavedModel serialize assert --- source/tests/consistent/io/test_io.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 574b3c472f..904576d4fe 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -160,8 +160,8 @@ def test_deep_eval(self) -> None: self.save_data_to_model(model_file, reference_data) deep_eval = DeepEval(model_file) self.assertIsInstance(deep_eval.get_model_def_script(), dict) - if not model_file.endswith(".savedmodel"): - # JAX SavedModel stores an executable graph, not a lossless model dict. + if not model_file.endswith((".savedmodel", ".savedmodeltf")): + # SavedModel formats store an executable graph, not a lossless model dict. serialized_data = self.get_data_from_model(model_file) np.testing.assert_equal(deep_eval.serialize(), serialized_data["model"]) if deep_eval.get_dim_fparam() > 0: