From 9227907addc1a4999bdac4f19baddd953399f958 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 08:00:56 +0800 Subject: [PATCH 1/8] fix(infer): normalize parameter shorthand before batching Standardize frame and atomic parameter shorthand in the common wrapper and backend entry points before automatic batching. Preserve full per-frame and per-atom arrays while broadcasting documented shared forms consistently across TensorFlow 2, PyTorch, JAX, and Paddle. Normalize PyTorch embedding extraction before split execution so eval_embedding, eval_descriptor, and eval_fitting_last_layer accept shared fparam, per-atom aparam, and scalar aparam forms. Cover two frames forced through one-frame backend batches. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/infer/deep_eval.py | 74 +++++++++---- deepmd/jax/infer/deep_eval.py | 9 ++ deepmd/pd/infer/deep_eval.py | 9 ++ deepmd/pt/infer/deep_eval.py | 20 ++++ deepmd/tf2/infer/deep_eval.py | 9 ++ .../test_deep_eval_parameter_shorthand.py | 101 ++++++++++++++++++ source/tests/consistent/io/test_io.py | 68 ++++++++++++ source/tests/pt/model/test_embedding.py | 95 ++++++++++++++++ 8 files changed, 364 insertions(+), 21 deletions(-) create mode 100644 source/tests/common/test_deep_eval_parameter_shorthand.py diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 05e40854d1..5e9f7de9b7 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -30,6 +30,48 @@ import ase.neighborlist +def _standardize_fparam_aparam( + fparam: np.ndarray | list | None, + aparam: np.ndarray | list | None, + nframes: int, + natoms: int, + dim_fparam: int, + dim_aparam: int, +) -> tuple[np.ndarray | None, np.ndarray | None]: + """Normalize documented parameter shorthand to frame-major arrays. + + This normalization must happen before automatic batching. In particular, + an ``(natoms, dim_aparam)`` shared atomic parameter has an atom axis first; + a batcher would otherwise mistake that axis for frames and slice it. + """ + if fparam is not None: + fparam = np.asarray(fparam) + if fparam.size == nframes * dim_fparam: + fparam = fparam.reshape(nframes, dim_fparam) + elif fparam.size == dim_fparam: + fparam = np.tile(fparam.reshape(1, dim_fparam), (nframes, 1)) + else: + raise RuntimeError( + "got wrong size of frame param, should be either " + f"{nframes} x {dim_fparam} or {dim_fparam}" + ) + if aparam is not None: + aparam = np.asarray(aparam) + if aparam.size == nframes * natoms * dim_aparam: + aparam = aparam.reshape(nframes, natoms, dim_aparam) + elif aparam.size == natoms * dim_aparam: + aparam = np.tile(aparam.reshape(1, natoms, dim_aparam), (nframes, 1, 1)) + elif aparam.size == dim_aparam: + aparam = np.tile(aparam.reshape(1, 1, dim_aparam), (nframes, natoms, 1)) + else: + raise RuntimeError( + "got wrong size of atomic param, should be either " + f"{nframes} x {natoms} x {dim_aparam} or " + f"{natoms} x {dim_aparam} or {dim_aparam}" + ) + return fparam, aparam + + class DeepEvalBackend(ABC): """Low-level Deep Evaluator interface. @@ -947,28 +989,18 @@ def _standard_input( coords = coords.reshape(nframes, natoms, 3) if cells is not None: cells = cells.reshape(nframes, 3, 3) - if fparam is not None: - fdim = self.get_dim_fparam() - if fparam.size == nframes * fdim: - fparam = np.reshape(fparam, [nframes, fdim]) - elif fparam.size == fdim: - fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) - else: - raise RuntimeError( - f"got wrong size of frame param, should be either {nframes} x {fdim} or {fdim}" - ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) if aparam is not None: - fdim = self.get_dim_aparam() - if aparam.size == nframes * natoms * fdim: - aparam = np.reshape(aparam, [nframes, natoms * fdim]) - elif aparam.size == natoms * fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) - elif aparam.size == fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) - else: - raise RuntimeError( - f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" - ) + # Preserve the historical flattened backend ABI used by the public + # wrapper; backend adapters normalize it back to frame-major 3-D. + aparam = aparam.reshape(nframes, natoms * self.get_dim_aparam()) return coords, cells, atom_types, fparam, aparam, nframes, natoms def get_sel_type(self) -> list[int]: diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 0e6c11ede6..9c4a151e81 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -35,6 +35,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -278,6 +279,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, charge_spin, request_defs diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index c8bb113495..5658a0e73c 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -28,6 +28,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -376,6 +377,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index dffdbee5e1..4f30448747 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -29,6 +29,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -544,6 +545,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( @@ -1312,6 +1321,17 @@ def eval_embedding( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + # Normalize shared parameter shorthand before auto batching. Otherwise + # a one-dimensional fparam/aparam is passed unchanged to every split, + # and _eval_embedding cannot reshape it to the split frame count. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) return self._eval_func(self._eval_embedding, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, charge_spin, dtype ) diff --git a/deepmd/tf2/infer/deep_eval.py b/deepmd/tf2/infer/deep_eval.py index af5037e2e6..35f78957af 100644 --- a/deepmd/tf2/infer/deep_eval.py +++ b/deepmd/tf2/infer/deep_eval.py @@ -32,6 +32,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -292,6 +293,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, request_defs diff --git a/source/tests/common/test_deep_eval_parameter_shorthand.py b/source/tests/common/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..aa54ca6622 --- /dev/null +++ b/source/tests/common/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for backend-level DeepEval parameter normalization.""" + +import numpy as np +import pytest + +from deepmd.infer.deep_eval import ( + _standardize_fparam_aparam, +) + + +NFRAMES = 3 +NATOMS = 4 +DIM_FPARAM = 2 +DIM_APARAM = 2 +FPARAM = np.array([0.25, -0.5], dtype=np.float64) +APARAM_PER_ATOM = np.arange(NATOMS * DIM_APARAM, dtype=np.float64).reshape( + NATOMS, DIM_APARAM +) +APARAM_ALL_ATOMS = np.array([0.3, -0.2], dtype=np.float64) + + +@pytest.mark.parametrize( + ("fparam", "expected"), + [ + (FPARAM.tolist(), np.tile(FPARAM, (NFRAMES, 1))), + ( + np.arange(NFRAMES * DIM_FPARAM).reshape(NFRAMES, DIM_FPARAM), + np.arange(NFRAMES * DIM_FPARAM).reshape(NFRAMES, DIM_FPARAM), + ), + ], + ids=("shared", "per-frame"), +) +def test_standardize_fparam(fparam, expected) -> None: + """Frame parameters become a canonical frame-major matrix.""" + actual, _ = _standardize_fparam_aparam( + fparam, + None, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("aparam", "expected"), + [ + ( + APARAM_PER_ATOM, + np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + ), + ( + APARAM_ALL_ATOMS.tolist(), + np.tile(APARAM_ALL_ATOMS, (NFRAMES, NATOMS, 1)), + ), + ( + np.arange(NFRAMES * NATOMS * DIM_APARAM).reshape( + NFRAMES, NATOMS, DIM_APARAM + ), + np.arange(NFRAMES * NATOMS * DIM_APARAM).reshape( + NFRAMES, NATOMS, DIM_APARAM + ), + ), + ], + ids=("shared-per-atom", "shared-all-atoms", "per-frame"), +) +def test_standardize_aparam(aparam, expected) -> None: + """Atomic shorthand is expanded before a batcher can slice its atom axis.""" + _, actual = _standardize_fparam_aparam( + None, + aparam, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("fparam", "aparam", "message"), + [ + (np.zeros(3), None, "wrong size of frame param"), + (None, np.zeros(3), "wrong size of atomic param"), + ], +) +def test_invalid_parameter_size_is_rejected(fparam, aparam, message) -> None: + """Report the documented contract instead of a backend reshape failure.""" + with pytest.raises(RuntimeError, match=message): + _standardize_fparam_aparam( + fparam, + aparam, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 2a34e2bbe5..6d8dc88571 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -39,6 +39,15 @@ class IOTest: # property model), skipped by the cross-backend round trips below. skip_backends: ClassVar[set[str]] = set() + def _has_fparam_aparam(self) -> bool: + """Whether the serialized fitting requires both parameter families.""" + fitting = self.data.get("model_def_script", {}).get("fitting_net", {}) + return ( + isinstance(fitting, dict) + and fitting.get("numb_fparam", 0) > 0 + and fitting.get("numb_aparam", 0) > 0 + ) + def get_data_from_model(self, model_file: str) -> dict: """Get data from a model file. @@ -156,6 +165,7 @@ def test_deep_eval(self) -> None: ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), ("tf2", 0) if DP_TEST_TF2_ONLY else (None, None), ("pytorch", 0), + ("paddle", 1) if self._has_fparam_aparam() else (None, None), ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), ): @@ -181,6 +191,10 @@ def test_deep_eval(self) -> None: aparam = np.ones((nframes, natoms, deep_eval.get_dim_aparam())) else: aparam = None + if backend_name in {"pytorch", "jax", "tf2", "paddle"} and ( + deep_eval.get_dim_fparam() > 0 and deep_eval.get_dim_aparam() > 0 + ): + self._assert_backend_parameter_shorthand(model_file, deep_eval) ret = deep_eval.eval( self.coords, self.box, @@ -238,6 +252,60 @@ def test_deep_eval(self) -> None: err_msg=f"backend {idx + 1} for rets_idx {rets_idx}", ) + def _assert_backend_parameter_shorthand( + self, model_file: str, deep_eval: DeepEval + ) -> None: + """Compare backend-direct shorthand with explicit frame-major inputs. + + Calling ``deep_eval.deep_eval`` deliberately bypasses the public + ``_standard_input`` normalization. A one-frame auto-batch size also + proves that shared per-atom parameters are expanded before the batcher + can mistake their atom axis for a frame axis. + """ + natoms = self.atype.shape[1] + nframes = 2 + coords = np.repeat(self.coords, nframes, axis=0) + boxes = np.repeat(self.box, nframes, axis=0) + atom_types = self.atype.reshape(-1) + fparam_shared = np.ones(deep_eval.get_dim_fparam()) + aparam_per_atom = np.ones((natoms, deep_eval.get_dim_aparam())) + fparam_full = np.tile(fparam_shared, (nframes, 1)) + aparam_full = np.tile(aparam_per_atom, (nframes, 1, 1)) + backend = DeepEval(model_file, auto_batch_size=natoms).deep_eval + + expected = backend.eval( + coords, + boxes, + atom_types, + fparam=fparam_full, + aparam=aparam_full, + ) + shorthand_cases = ( + (fparam_shared.tolist(), aparam_per_atom), + ( + fparam_shared, + np.ones(deep_eval.get_dim_aparam()), + ), + ) + for fparam, aparam in shorthand_cases: + actual = backend.eval( + coords, + boxes, + atom_types, + fparam=fparam, + aparam=aparam, + ) + self.assertEqual(actual.keys(), expected.keys()) + for name in actual: + np.testing.assert_allclose( + actual[name], + expected[name], + rtol=1e-12, + atol=1e-12, + equal_nan=True, + err_msg=f"backend-direct shorthand output {name}", + ) + class TestDeepPot(unittest.TestCase, IOTest): def setUp(self) -> None: diff --git a/source/tests/pt/model/test_embedding.py b/source/tests/pt/model/test_embedding.py index 2617a0afba..f2afcf4c6b 100644 --- a/source/tests/pt/model/test_embedding.py +++ b/source/tests/pt/model/test_embedding.py @@ -314,6 +314,101 @@ def test_eval_embedding_dtype_fp64(self) -> None: np.float64, ) + def test_backend_embedding_normalizes_parameter_shorthand_before_batching( + self, + ) -> None: + """Shared fparam/aparam forms must survive one-frame auto batching.""" + params = _se_e2_a_params() + params["fitting_net"]["numb_fparam"] = 2 + params["fitting_net"]["numb_aparam"] = 1 + model = get_model(params) + _randomize(model) + path = self._save_checkpoint(model, params, "se_e2_a_params.pt") + + natoms = int(self.atype_np.shape[0]) + coords = np.repeat(self.coord_np, 2, axis=0) + cells = np.repeat(self.cell_np, 2, axis=0) + fparam_shared = np.array([0.25, -0.5], dtype=np.float64) + fparam_full = np.repeat(fparam_shared[None, :], 2, axis=0) + aparam_shared = np.linspace(0.1, 0.7, natoms, dtype=np.float64)[:, None] + aparam_full = np.repeat(aparam_shared[None, :, :], 2, axis=0) + + # A batch budget of exactly natoms forces each of the two frames into + # a separate backend call and exposes normalization done too late. + dp = DeepPot(path, auto_batch_size=natoms, no_jit=True) + backend = dp.deep_eval + self.assertIsInstance(backend, PTDeepEval) + self.assertEqual(backend.auto_batch_size.current_batch_size, natoms) + + full = backend.eval_embedding( + coords, + cells, + self.atype_np, + fparam=fparam_full, + aparam=aparam_full, + ) + shared = backend.eval_embedding( + coords, + cells, + self.atype_np, + fparam=fparam_shared, + aparam=aparam_shared, + ) + for full_value, shared_value in zip(full, shared, strict=True): + np.testing.assert_allclose(shared_value, full_value) + + np.testing.assert_allclose( + backend.eval_descriptor( + coords, + cells, + self.atype_np, + fparam=fparam_shared, + aparam=aparam_shared, + ), + backend.eval_descriptor( + coords, + cells, + self.atype_np, + fparam=fparam_full, + aparam=aparam_full, + ), + ) + np.testing.assert_allclose( + backend.eval_fitting_last_layer( + coords, + cells, + self.atype_np, + fparam=fparam_shared, + aparam=aparam_shared, + ), + backend.eval_fitting_last_layer( + coords, + cells, + self.atype_np, + fparam=fparam_full, + aparam=aparam_full, + ), + ) + + scalar_aparam = np.array([0.35], dtype=np.float64) + scalar_full = np.full((2, natoms, 1), scalar_aparam.item()) + scalar = backend.eval_embedding( + coords, + cells, + self.atype_np, + fparam=fparam_shared, + aparam=scalar_aparam, + ) + scalar_reference = backend.eval_embedding( + coords, + cells, + self.atype_np, + fparam=fparam_full, + aparam=scalar_full, + ) + for scalar_value, reference_value in zip(scalar, scalar_reference, strict=True): + np.testing.assert_allclose(scalar_value, reference_value) + def test_legacy_frozen_model_uses_baked_in_hook(self) -> None: # Frozen ``.pth`` files predating ``forward_embedding`` still carry the # descriptor / fitting hooks baked into the TorchScript module. The From bbfc000ac2b1b4c39f7ab242033342f6a50869b6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:58:19 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/test_deep_eval_parameter_shorthand.py | 1 - 1 file changed, 1 deletion(-) diff --git a/source/tests/common/test_deep_eval_parameter_shorthand.py b/source/tests/common/test_deep_eval_parameter_shorthand.py index aa54ca6622..4aa18fe93a 100644 --- a/source/tests/common/test_deep_eval_parameter_shorthand.py +++ b/source/tests/common/test_deep_eval_parameter_shorthand.py @@ -8,7 +8,6 @@ _standardize_fparam_aparam, ) - NFRAMES = 3 NATOMS = 4 DIM_FPARAM = 2 From 8d80b8579fec8e5f5ed66ca599336fe53b9979e9 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 15:24:44 +0800 Subject: [PATCH 3/8] test(io): drop the unusable paddle round trip deserialize_to_file only writes .json, serialize_from_file raises NotImplementedError, and the .json reader rejects fparam/aparam, so the paddle entry could not complete the round trip. Its normalization stays covered by test_deep_eval_parameter_shorthand.py. --- source/tests/consistent/io/test_io.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 6d8dc88571..cd13b78232 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -39,15 +39,6 @@ class IOTest: # property model), skipped by the cross-backend round trips below. skip_backends: ClassVar[set[str]] = set() - def _has_fparam_aparam(self) -> bool: - """Whether the serialized fitting requires both parameter families.""" - fitting = self.data.get("model_def_script", {}).get("fitting_net", {}) - return ( - isinstance(fitting, dict) - and fitting.get("numb_fparam", 0) > 0 - and fitting.get("numb_aparam", 0) > 0 - ) - def get_data_from_model(self, model_file: str) -> dict: """Get data from a model file. @@ -165,7 +156,6 @@ def test_deep_eval(self) -> None: ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), ("tf2", 0) if DP_TEST_TF2_ONLY else (None, None), ("pytorch", 0), - ("paddle", 1) if self._has_fparam_aparam() else (None, None), ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), ): @@ -191,7 +181,11 @@ def test_deep_eval(self) -> None: aparam = np.ones((nframes, natoms, deep_eval.get_dim_aparam())) else: aparam = None - if backend_name in {"pytorch", "jax", "tf2", "paddle"} and ( + # Paddle is absent from the loop above: deserialize_to_file only + # writes .json, serialize_from_file is not implemented, and the + # .json reader rejects fparam/aparam. Its normalization is covered + # by source/tests/common/test_deep_eval_parameter_shorthand.py. + if backend_name in {"pytorch", "jax", "tf2"} and ( deep_eval.get_dim_fparam() > 0 and deep_eval.get_dim_aparam() > 0 ): self._assert_backend_parameter_shorthand(model_file, deep_eval) From 9867233019b463ba89236b9d237af4894c8f9328 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 23:16:17 +0800 Subject: [PATCH 4/8] fix(infer): complete parameter shorthand coverage Normalize parameter shorthand in the remaining dpmodel, pt_expt, and TensorFlow adapters before automatic batching. Strengthen the regressions with distinct frames and parameters, fresh batchers, and direct Paddle and pt_expt adapter coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/infer/deep_eval.py | 9 ++ deepmd/infer/deep_eval.py | 27 ++++ deepmd/pt_expt/infer/deep_eval.py | 20 ++- deepmd/tf/infer/deep_eval.py | 9 ++ source/tests/consistent/io/test_io.py | 81 +++++++---- .../pd/test_deep_eval_parameter_shorthand.py | 54 ++++++++ source/tests/pt/model/test_embedding.py | 126 +++++++++--------- .../pt_expt/infer/test_parameter_shorthand.py | 52 ++++++++ 8 files changed, 282 insertions(+), 96 deletions(-) create mode 100644 source/tests/pd/test_deep_eval_parameter_shorthand.py create mode 100644 source/tests/pt_expt/infer/test_parameter_shorthand.py diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index e86322866b..bfa4a80d96 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -40,6 +40,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -240,6 +241,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, request_defs diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 58cde679aa..2621607cb7 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -43,6 +43,33 @@ def _standardize_fparam_aparam( This normalization must happen before automatic batching. In particular, an ``(natoms, dim_aparam)`` shared atomic parameter has an atom axis first; a batcher would otherwise mistake that axis for frames and slice it. + + Parameters + ---------- + fparam : np.ndarray or list or None + Frame parameters in full ``(nframes, dim_fparam)`` form or shared + ``(dim_fparam,)`` shorthand. + aparam : np.ndarray or list or None + Atomic parameters in full ``(nframes, natoms, dim_aparam)`` form, + shared-per-atom ``(natoms, dim_aparam)`` shorthand, or shared-for-all + ``(dim_aparam,)`` shorthand. + nframes : int + Number of input frames. + natoms : int + Number of atoms in each frame. + dim_fparam : int + Number of frame-parameter components. + dim_aparam : int + Number of atomic-parameter components. + + Returns + ------- + normalized_fparam : np.ndarray or None + Frame parameters with shape ``(nframes, dim_fparam)`` when provided. + normalized_aparam : np.ndarray or None + Atomic parameters with shape ``(nframes, natoms, dim_aparam)`` when + provided. The public wrapper may subsequently flatten the last two + axes to preserve its historical backend ABI. """ if fparam is not None: fparam = np.asarray(fparam) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1b97ef7f50..d2b24bfccc 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -41,6 +41,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -848,10 +849,15 @@ def eval( Calculate the atomic energy and virial fparam The frame parameter. - The array should be of size nframes x dim_fparam. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. aparam - The atomic parameter. - The array should be of size nframes x natoms x dim_aparam. + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. charge_spin The charge and spin values for each frame. The array should be reshape-compatible with nframes x 2, where the first @@ -875,6 +881,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) spins = kwargs.get("spin") if self._is_spin and spins is None: diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 30107cb693..3d386c48dc 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -32,6 +32,7 @@ ) from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -761,6 +762,14 @@ def eval( coords, atom_types, ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) output = self._eval_func(self._eval_inner, numb_test, natoms)( coords, cells, diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index cd13b78232..9c70217ec0 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -181,13 +181,10 @@ def test_deep_eval(self) -> None: aparam = np.ones((nframes, natoms, deep_eval.get_dim_aparam())) else: aparam = None - # Paddle is absent from the loop above: deserialize_to_file only - # writes .json, serialize_from_file is not implemented, and the - # .json reader rejects fparam/aparam. Its normalization is covered - # by source/tests/common/test_deep_eval_parameter_shorthand.py. - if backend_name in {"pytorch", "jax", "tf2"} and ( - deep_eval.get_dim_fparam() > 0 and deep_eval.get_dim_aparam() > 0 - ): + # Paddle is absent from the loop above because its frozen .json + # reader rejects fparam/aparam. A dedicated adapter test executes + # its normalization path without relying on that frozen format. + if deep_eval.get_dim_fparam() > 0 and deep_eval.get_dim_aparam() > 0: self._assert_backend_parameter_shorthand(model_file, deep_eval) ret = deep_eval.eval( self.coords, @@ -258,30 +255,50 @@ def _assert_backend_parameter_shorthand( """ natoms = self.atype.shape[1] nframes = 2 - coords = np.repeat(self.coords, nframes, axis=0) + coords = np.concatenate((self.coords, self.coords + 0.125), axis=0) boxes = np.repeat(self.box, nframes, axis=0) - atom_types = self.atype.reshape(-1) - fparam_shared = np.ones(deep_eval.get_dim_fparam()) - aparam_per_atom = np.ones((natoms, deep_eval.get_dim_aparam())) + boxes[1, 0] += 0.25 + atom_types = np.repeat(self.atype, nframes, axis=0) + dim_fparam = deep_eval.get_dim_fparam() + dim_aparam = deep_eval.get_dim_aparam() + fparam_shared = np.linspace(0.25, -0.5, dim_fparam) + aparam_per_atom = np.linspace(0.1, 0.9, natoms * dim_aparam).reshape( + natoms, dim_aparam + ) fparam_full = np.tile(fparam_shared, (nframes, 1)) aparam_full = np.tile(aparam_per_atom, (nframes, 1, 1)) - backend = DeepEval(model_file, auto_batch_size=natoms).deep_eval - - expected = backend.eval( - coords, - boxes, - atom_types, - fparam=fparam_full, - aparam=aparam_full, - ) - shorthand_cases = ( - (fparam_shared.tolist(), aparam_per_atom), + fparam_per_frame = np.stack((fparam_shared, fparam_shared + 0.375)) + aparam_per_frame = np.stack((aparam_per_atom, aparam_per_atom + 0.2)) + aparam_all_atoms = np.linspace(-0.3, 0.4, dim_aparam) + + cases = ( ( - fparam_shared, - np.ones(deep_eval.get_dim_aparam()), + fparam_shared.tolist(), + aparam_per_atom, + fparam_full, + aparam_full, + "shared-frame-and-per-atom", + ), + ( + fparam_per_frame, + aparam_all_atoms, + fparam_per_frame, + np.tile(aparam_all_atoms, (nframes, natoms, 1)), + "shared-all-atoms", + ), + ( + fparam_per_frame, + aparam_per_frame, + fparam_per_frame, + aparam_per_frame, + "full-frame-major", ), ) - for fparam, aparam in shorthand_cases: + for fparam, aparam, full_fparam, full_aparam, case_name in cases: + # Run the shorthand/full input first on a fresh auto-batcher so a + # GPU runner cannot grow the batch size on a preceding reference + # call and silently stop exercising the split path. + backend = DeepEval(model_file, auto_batch_size=natoms).deep_eval actual = backend.eval( coords, boxes, @@ -289,15 +306,23 @@ def _assert_backend_parameter_shorthand( fparam=fparam, aparam=aparam, ) + reference_backend = DeepEval(model_file, auto_batch_size=False).deep_eval + expected = reference_backend.eval( + coords, + boxes, + atom_types, + fparam=full_fparam, + aparam=full_aparam, + ) self.assertEqual(actual.keys(), expected.keys()) for name in actual: np.testing.assert_allclose( actual[name], expected[name], - rtol=1e-12, - atol=1e-12, + rtol=1e-10, + atol=1e-10, equal_nan=True, - err_msg=f"backend-direct shorthand output {name}", + err_msg=f"backend-direct {case_name} output {name}", ) diff --git a/source/tests/pd/test_deep_eval_parameter_shorthand.py b/source/tests/pd/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..b3c065ae97 --- /dev/null +++ b/source/tests/pd/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exercise parameter shorthand through the Paddle DeepEval adapter.""" + +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + MagicMock, +) + +import numpy as np +import pytest + +pytest.importorskip("paddle") + +from deepmd.pd.infer.deep_eval import ( + DeepEval, +) + + +def test_eval_standardizes_parameter_shorthand_before_dispatch() -> None: + """The Paddle adapter must pass frame-major parameters to its evaluator.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._get_request_defs = lambda atomic: [SimpleNamespace(name="energy")] + backend._eval_func = lambda inner, numb_test, numb_atoms: inner + backend._eval_model = MagicMock(return_value=(np.zeros((nframes, 1)),)) + + result = backend.eval( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._eval_model.call_args.args[3], np.tile(fparam, (nframes, 1)) + ) + np.testing.assert_array_equal( + backend._eval_model.call_args.args[4], np.tile(aparam, (nframes, 1, 1)) + ) + assert result.keys() == {"energy"} diff --git a/source/tests/pt/model/test_embedding.py b/source/tests/pt/model/test_embedding.py index f2afcf4c6b..afc9a827d0 100644 --- a/source/tests/pt/model/test_embedding.py +++ b/source/tests/pt/model/test_embedding.py @@ -326,88 +326,84 @@ def test_backend_embedding_normalizes_parameter_shorthand_before_batching( path = self._save_checkpoint(model, params, "se_e2_a_params.pt") natoms = int(self.atype_np.shape[0]) - coords = np.repeat(self.coord_np, 2, axis=0) + coords = np.concatenate((self.coord_np, self.coord_np + 0.125), axis=0) cells = np.repeat(self.cell_np, 2, axis=0) + cells[1, 0] += 0.25 fparam_shared = np.array([0.25, -0.5], dtype=np.float64) fparam_full = np.repeat(fparam_shared[None, :], 2, axis=0) aparam_shared = np.linspace(0.1, 0.7, natoms, dtype=np.float64)[:, None] aparam_full = np.repeat(aparam_shared[None, :, :], 2, axis=0) - # A batch budget of exactly natoms forces each of the two frames into - # a separate backend call and exposes normalization done too late. - dp = DeepPot(path, auto_batch_size=natoms, no_jit=True) - backend = dp.deep_eval - self.assertIsInstance(backend, PTDeepEval) - self.assertEqual(backend.auto_batch_size.current_batch_size, natoms) - - full = backend.eval_embedding( - coords, - cells, - self.atype_np, - fparam=fparam_full, - aparam=aparam_full, - ) - shared = backend.eval_embedding( - coords, - cells, - self.atype_np, - fparam=fparam_shared, - aparam=aparam_shared, - ) - for full_value, shared_value in zip(full, shared, strict=True): - np.testing.assert_allclose(shared_value, full_value) - - np.testing.assert_allclose( - backend.eval_descriptor( - coords, - cells, - self.atype_np, - fparam=fparam_shared, - aparam=aparam_shared, - ), - backend.eval_descriptor( + def assert_method_matches_full( + method_name: str, + fparam: np.ndarray, + aparam: np.ndarray, + full_fparam: np.ndarray, + full_aparam: np.ndarray, + ) -> None: + # A fresh batch budget of exactly natoms forces the shorthand call + # to split the two frames even on GPU, where successful calls grow + # the auto-batch size. The reference is intentionally unbatched. + backend = DeepPot(path, auto_batch_size=natoms, no_jit=True).deep_eval + self.assertIsInstance(backend, PTDeepEval) + self.assertEqual(backend.auto_batch_size.current_batch_size, natoms) + actual = getattr(backend, method_name)( coords, cells, self.atype_np, - fparam=fparam_full, - aparam=aparam_full, - ), - ) - np.testing.assert_allclose( - backend.eval_fitting_last_layer( - coords, - cells, - self.atype_np, - fparam=fparam_shared, - aparam=aparam_shared, - ), - backend.eval_fitting_last_layer( + fparam=fparam, + aparam=aparam, + ) + + reference_backend = DeepPot( + path, auto_batch_size=False, no_jit=True + ).deep_eval + expected = getattr(reference_backend, method_name)( coords, cells, self.atype_np, - fparam=fparam_full, - aparam=aparam_full, - ), - ) + fparam=full_fparam, + aparam=full_aparam, + ) + actual_values = actual if isinstance(actual, tuple) else (actual,) + expected_values = expected if isinstance(expected, tuple) else (expected,) + for actual_value, expected_value in zip( + actual_values, expected_values, strict=True + ): + np.testing.assert_allclose(actual_value, expected_value) + + for method_name in ( + "eval_embedding", + "eval_descriptor", + "eval_fitting_last_layer", + ): + assert_method_matches_full( + method_name, + fparam_shared, + aparam_shared, + fparam_full, + aparam_full, + ) scalar_aparam = np.array([0.35], dtype=np.float64) scalar_full = np.full((2, natoms, 1), scalar_aparam.item()) - scalar = backend.eval_embedding( - coords, - cells, - self.atype_np, - fparam=fparam_shared, - aparam=scalar_aparam, + assert_method_matches_full( + "eval_embedding", + fparam_shared, + scalar_aparam, + fparam_full, + scalar_full, ) - scalar_reference = backend.eval_embedding( - coords, - cells, - self.atype_np, - fparam=fparam_full, - aparam=scalar_full, + + # Distinct frame-major values ensure batching preserves the frame and + # atom axes instead of merely producing a shape-compatible result. + assert_method_matches_full( + "eval_embedding", + np.stack((fparam_shared, fparam_shared + 0.375)), + np.stack((aparam_shared, aparam_shared + 0.2)), + np.stack((fparam_shared, fparam_shared + 0.375)), + np.stack((aparam_shared, aparam_shared + 0.2)), ) - for scalar_value, reference_value in zip(scalar, scalar_reference, strict=True): - np.testing.assert_allclose(scalar_value, reference_value) def test_legacy_frozen_model_uses_baked_in_hook(self) -> None: # Frozen ``.pth`` files predating ``forward_embedding`` still carry the diff --git a/source/tests/pt_expt/infer/test_parameter_shorthand.py b/source/tests/pt_expt/infer/test_parameter_shorthand.py new file mode 100644 index 0000000000..1128014a0f --- /dev/null +++ b/source/tests/pt_expt/infer/test_parameter_shorthand.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exercise parameter shorthand through the pt_expt DeepEval adapter.""" + +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + MagicMock, +) + +import numpy as np + +from deepmd.pt_expt.infer.deep_eval import ( + DeepEval, +) + + +def test_eval_standardizes_parameter_shorthand_before_dispatch() -> None: + """The pt_expt adapter must normalize parameters before auto batching.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend._is_spin = False + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._get_request_defs = lambda atomic: [SimpleNamespace(name="energy")] + backend._eval_func = lambda inner, numb_test, numb_atoms: inner + backend._eval_model = MagicMock(return_value=(np.zeros((nframes, 1)),)) + + result = backend.eval( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._eval_model.call_args.args[3], np.tile(fparam, (nframes, 1)) + ) + np.testing.assert_array_equal( + backend._eval_model.call_args.args[4], np.tile(aparam, (nframes, 1, 1)) + ) + assert result.keys() == {"energy"} From f785c9e0bad5c95e4bf68c5c02fce58ced90b9d4 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:53:08 +0800 Subject: [PATCH 5/8] fix(tf): normalize descriptor parameters before batching Coding-Agent: Codex\nCodex-Version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh --- deepmd/tf/infer/deep_eval.py | 11 ++++ .../tf/test_deep_eval_parameter_shorthand.py | 55 +++++++++++++++++++ 2 files changed, 66 insertions(+) create mode 100644 source/tests/tf/test_deep_eval_parameter_shorthand.py diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 3d386c48dc..0209bc694f 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1093,6 +1093,17 @@ def eval_descriptor( coords, atom_types, ) + # Canonicalize shared parameter shorthands before AutoBatchSize slices + # the frame axis; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)( coords, cells, diff --git a/source/tests/tf/test_deep_eval_parameter_shorthand.py b/source/tests/tf/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..1b359f5949 --- /dev/null +++ b/source/tests/tf/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for TensorFlow DeepEval parameter shorthand handling.""" + +import numpy as np + +from deepmd.tf.infer.deep_eval import ( + DeepEval, +) + + +class _CapturingAutoBatchSize: + """Capture inputs at the auto-batching boundary without loading a model.""" + + def __init__(self) -> None: + self.kwargs: dict | None = None + + def execute_all( + self, + inner_func, + nframes: int, + natoms: int, + *args, + **kwargs, + ) -> np.ndarray: + """Record canonical inputs that would be sliced into frame batches.""" + self.kwargs = kwargs + return np.zeros((nframes, natoms, 1), dtype=np.float64) + + +def test_eval_descriptor_normalizes_shared_parameters_before_batching() -> None: + """Shared parameters must gain a frame axis before auto batching.""" + auto_batch_size = _CapturingAutoBatchSize() + backend = object.__new__(DeepEval) + backend.dfparam = 2 + backend.daparam = 1 + backend.auto_batch_size = auto_batch_size + + fparam = np.array([0.25, -0.5]) + aparam = np.array([[0.1], [0.2]]) + descriptor = backend.eval_descriptor( + np.zeros((2, 2, 3)), + None, + np.array([0, 1]), + fparam=fparam, + aparam=aparam, + ) + + assert descriptor.shape == (2, 2, 1) + assert auto_batch_size.kwargs is not None + np.testing.assert_array_equal( + auto_batch_size.kwargs["fparam"], np.tile(fparam, (2, 1)) + ) + np.testing.assert_array_equal( + auto_batch_size.kwargs["aparam"], np.tile(aparam[None, :, :], (2, 1, 1)) + ) From 866ce34bfd5fddb0fd301a172b1d69ac5bcd6241 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 22:06:47 +0800 Subject: [PATCH 6/8] fix(pt-expt): normalize shorthand in descriptor and fitting-last-layer routes Wire _standardize_fparam_aparam into eval_descriptor and eval_fitting_last_layer so shared per-atom shorthand cannot be mistaken for a frame axis before _prepare_nlist_inputs reshapes it. Extend the pt_expt adapter regressions to cover both routes. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- deepmd/pt_expt/infer/deep_eval.py | 26 +++++ .../pt_expt/infer/test_parameter_shorthand.py | 107 ++++++++++++++++++ 2 files changed, 133 insertions(+) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index d2b24bfccc..3daf55b825 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -2217,6 +2217,19 @@ def eval_descriptor( "eval_descriptor is not supported for this model type " f"({type(self._dpmodel).__name__})." ) + nframes = coords.shape[0] + natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1] + # Canonicalize shared parameter shorthands before _prepare_nlist_inputs + # reshapes them; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) ( ext_coord_t, ext_atype_t, @@ -2288,6 +2301,19 @@ def eval_fitting_last_layer( "eval_fitting_last_layer is not supported for this model type " f"({type(self._dpmodel).__name__})." ) + nframes = coords.shape[0] + natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1] + # Canonicalize shared parameter shorthands before _prepare_nlist_inputs + # reshapes them; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) ( ext_coord_t, ext_atype_t, diff --git a/source/tests/pt_expt/infer/test_parameter_shorthand.py b/source/tests/pt_expt/infer/test_parameter_shorthand.py index 1128014a0f..8be4168cb1 100644 --- a/source/tests/pt_expt/infer/test_parameter_shorthand.py +++ b/source/tests/pt_expt/infer/test_parameter_shorthand.py @@ -50,3 +50,110 @@ def test_eval_standardizes_parameter_shorthand_before_dispatch() -> None: backend._eval_model.call_args.args[4], np.tile(aparam, (nframes, 1, 1)) ) assert result.keys() == {"energy"} + + +def _new_backend() -> DeepEval: + """Build an uninitialized pt_expt DeepEval for adapter mocking.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._require_dpmodel = lambda name: None + backend._is_spin_model = lambda: False + return backend + + +def test_eval_descriptor_standardizes_parameter_shorthand() -> None: + """The descriptor route must normalize parameters before reshaping.""" + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend = _new_backend() + dp_am = MagicMock() + backend._dpmodel = MagicMock() + backend._dpmodel.get_dp_atomic_model.return_value = dp_am + descriptor = MagicMock() + descriptor.detach.return_value = descriptor + descriptor.cpu.return_value = descriptor + descriptor.numpy.return_value = np.zeros((nframes, natoms, 4)) + dp_am.descriptor.return_value = (descriptor,) + backend._prepare_nlist_inputs = MagicMock( + return_value=(None, None, None, None, None, None, None, nframes, natoms) + ) + + backend.eval_descriptor( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[3], + np.tile(fparam, (nframes, 1)), + ) + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[4], + np.tile(aparam, (nframes, 1, 1)), + ) + + +def test_eval_fitting_last_layer_standardizes_parameter_shorthand() -> None: + """The fitting-last-layer route must normalize parameters before reshaping.""" + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend = _new_backend() + dp_am = MagicMock() + backend._dpmodel = MagicMock() + backend._dpmodel.get_dp_atomic_model.return_value = dp_am + dp_am.descriptor.return_value = ( + np.zeros((nframes, natoms, 4)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3)), + ) + out = MagicMock() + out.detach.return_value = out + out.cpu.return_value = out + out.numpy.return_value = np.zeros((nframes, natoms, 4)) + dp_am.fitting_net = MagicMock(return_value={"middle_output": out}) + ext_atype_t = np.zeros((nframes, natoms), dtype=np.int64) + backend._prepare_nlist_inputs = MagicMock( + return_value=( + None, + ext_atype_t, + None, + None, + None, + None, + None, + nframes, + natoms, + ) + ) + + backend.eval_fitting_last_layer( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[3], + np.tile(fparam, (nframes, 1)), + ) + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[4], + np.tile(aparam, (nframes, 1, 1)), + ) From 1709902d73dc6dfb92795b0cb956d3716d9c5df0 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 22:13:53 +0800 Subject: [PATCH 7/8] test(io): assert parameter shorthand preserves per-frame variation The backend-direct shorthand comparison never checked that the two input frames actually differ, so a degenerate expansion that reuses one frame's parameters for every frame would also make the frame-major reference degenerate and slip past the allclose comparison. Assert the per-frame energy values differ on every shorthand case. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- source/tests/consistent/io/test_io.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 9c70217ec0..662f9243cb 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -315,6 +315,20 @@ def _assert_backend_parameter_shorthand( aparam=full_aparam, ) self.assertEqual(actual.keys(), expected.keys()) + # The frame-major expansion must preserve per-frame parameter + # variation. Without this, a degenerate expansion that reuses one + # frame's parameters for every frame would also make the reference + # degenerate, so the allclose comparison alone could not catch it. + frame_key = next( + (key for key in ("energy_redu", "energy") if key in actual), None + ) + self.assertIsNotNone(frame_key) + frame_values = actual[frame_key] + self.assertEqual(frame_values.shape[0], 2) + self.assertFalse( + np.allclose(frame_values[0], frame_values[1]), + msg=f"backend-direct {case_name} params collapsed across frames", + ) for name in actual: np.testing.assert_allclose( actual[name], From 692789dfb3a19acb7b5346f172bceafd1bf552ec Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 23:55:32 +0800 Subject: [PATCH 8/8] test(dpmodel): cover parameter shorthand batching Stack the dpmodel-specific regression on the shared backend normalization from #5857. Use distinct coordinates and frame-major parameters, run shorthand inputs on fresh batchers, and verify split evaluation preserves frame order and distinct energies. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../infer/test_dpmodel_deep_eval_params.py | 226 ++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 source/tests/infer/test_dpmodel_deep_eval_params.py diff --git a/source/tests/infer/test_dpmodel_deep_eval_params.py b/source/tests/infer/test_dpmodel_deep_eval_params.py new file mode 100644 index 0000000000..321089caa6 --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_params.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Direct-backend tests for dpmodel DeepEval parameter shorthand.""" + +import copy + +import numpy as np +import pytest + +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.dpmodel.utils.serialization import ( + save_dp_model, +) +from deepmd.infer import ( + DeepEval, +) + +MODEL_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.5, + "rcut": 6.0, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 2, + }, +} +ATOM_TYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) +COORD = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=np.float64, +).reshape(len(ATOM_TYPES), 3) +BOX = np.diag([13.0, 13.0, 13.0]) +NFRAMES = 2 +COORDS = np.stack((COORD, COORD + 0.125)) +BOXES = np.tile(BOX, (NFRAMES, 1, 1)) +BOXES[1, 0, 0] += 0.25 +FPARAM = np.array([0.25, -0.5], dtype=np.float64) +APARAM_PER_ATOM = ( + np.arange(len(ATOM_TYPES) * 2, dtype=np.float64).reshape(len(ATOM_TYPES), 2) / 10.0 +) +APARAM_GLOBAL = np.array([0.3, -0.2], dtype=np.float64) + + +@pytest.fixture(scope="module") +def model_file(tmp_path_factory: pytest.TempPathFactory) -> str: + """Serialize a two-dimensional fparam/aparam model for direct inference.""" + model = get_model(copy.deepcopy(MODEL_CONFIG)) + path = tmp_path_factory.mktemp("dpmodel_params") / "model.dp" + save_dp_model( + str(path), + { + "model": model.serialize(), + "model_def_script": MODEL_CONFIG, + "backend": "dpmodel", + }, + ) + return str(path) + + +def _backend(model_file: str, auto_batch_size: bool | int): + """Return the low-level backend without public input normalization.""" + return DeepEval(model_file, auto_batch_size=auto_batch_size).deep_eval + + +def _assert_outputs_equal(actual: dict, expected: dict) -> None: + assert actual.keys() == expected.keys() + for name in actual: + np.testing.assert_allclose( + actual[name], + expected[name], + rtol=1e-10, + atol=1e-10, + equal_nan=True, + ) + + +@pytest.mark.parametrize("auto_batch_size", [False, len(ATOM_TYPES)]) +def test_shared_fparam_is_tiled_before_batching( + model_file: str, auto_batch_size: bool | int +) -> None: + """A single frame-parameter vector applies to every input frame.""" + full_aparam = np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)) + # Evaluate shorthand first on a fresh batcher. On GPU, a successful + # reference call would otherwise grow the batch size and remove the split + # this regression is intended to exercise. + actual = _backend(model_file, auto_batch_size).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=FPARAM.tolist(), + aparam=full_aparam, + ) + expected = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=np.tile(FPARAM, (NFRAMES, 1)), + aparam=full_aparam, + ) + + _assert_outputs_equal(actual, expected) + + +@pytest.mark.parametrize("auto_batch_size", [False, len(ATOM_TYPES)]) +@pytest.mark.parametrize( + ("shared_aparam", "full_aparam"), + [ + ( + APARAM_PER_ATOM, + np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + ), + ( + APARAM_GLOBAL, + np.tile(APARAM_GLOBAL, (NFRAMES, len(ATOM_TYPES), 1)), + ), + ], + ids=("per-atom", "all-atoms"), +) +def test_shared_aparam_is_tiled_before_batching( + model_file: str, + auto_batch_size: bool | int, + shared_aparam: np.ndarray, + full_aparam: np.ndarray, +) -> None: + """Both documented atomic-parameter shorthand forms are frame-major.""" + full_fparam = np.tile(FPARAM, (NFRAMES, 1)) + actual = _backend(model_file, auto_batch_size).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=shared_aparam, + ) + expected = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=full_aparam, + ) + + _assert_outputs_equal(actual, expected) + + +def test_distinct_full_parameters_preserve_frame_order_when_split( + model_file: str, +) -> None: + """Split evaluation preserves distinct frame-major parameter rows.""" + fparam = np.stack((FPARAM, FPARAM + np.array([0.4, -0.2]))) + aparam = np.stack((APARAM_PER_ATOM, APARAM_PER_ATOM + 0.35)) + + split = _backend(model_file, auto_batch_size=len(ATOM_TYPES)).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=fparam, + aparam=aparam, + ) + unsplit = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=fparam, + aparam=aparam, + ) + + _assert_outputs_equal(split, unsplit) + assert not np.allclose(split["energy_redu"][0], split["energy_redu"][1]) + + +@pytest.mark.parametrize( + ("parameter", "value", "message"), + [ + ("fparam", np.zeros(3), "wrong size of frame param"), + ("aparam", np.zeros(3), "wrong size of atomic param"), + ], +) +def test_invalid_parameter_size_has_clear_error( + model_file: str, parameter: str, value: np.ndarray, message: str +) -> None: + """Reject invalid sizes before NumPy emits an opaque reshape error.""" + evaluator = _backend(model_file, auto_batch_size=False) + parameters = { + "fparam": np.tile(FPARAM, (NFRAMES, 1)), + "aparam": np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + } + parameters[parameter] = value + + with pytest.raises(RuntimeError, match=message): + evaluator.eval(COORDS, BOXES, ATOM_TYPES, **parameters)