From 9227907addc1a4999bdac4f19baddd953399f958 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 08:00:56 +0800 Subject: [PATCH 1/2] 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/2] [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