fix(infer): normalize parameter shorthand before batching #5857
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
Collaborator
There was a problem hiding this comment. The docstring explains the rationale well but skips the numpydoc With six parameters and a two-element tuple return, and given this is now the single definition of the shorthand contract for four backends, it is worth documenting each argument and the returned shapes explicitly - in particular that
Contributor
Author
There was a problem hiding this comment. Fixed in 9867233. The shared helper now has complete numpydoc Parameters and Returns sections, including the canonical 2-D fparam and 3-D aparam return shapes and the public wrapper flattening note. Coding agent: Codex |
||
|
|
||
| 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. | ||
|
Comment on lines
+43
to
+45
Contributor
Author
There was a problem hiding this comment. Verified across the public wrapper and the dpmodel, JAX, Paddle, PyTorch, pt_expt, TensorFlow v1, and TF2 adapters: normalizing to a frame-major representation at this boundary prevents |
||
| """ | ||
| 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. | ||
|
|
||
|
|
@@ -948,28 +990,18 @@ def _standard_input( | |
| coords = coords.reshape(nframes, natoms, 3) | ||
| if cells is not None: | ||
| cells = cells.reshape(nframes, 3, 3) | ||
| fparam, aparam = _standardize_fparam_aparam( | ||
| fparam, | ||
| aparam, | ||
| nframes, | ||
| natoms, | ||
| self.get_dim_fparam(), | ||
| self.get_dim_aparam(), | ||
| ) | ||
| if aparam is not None: | ||
| # 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]: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| # 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, | ||
| ) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -181,6 +181,14 @@ 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 | ||
|
Collaborator
There was a problem hiding this comment. This comment claims more than the referenced file delivers. After Either restore some executing coverage for pd, or reword both the comment and the PR body to say plainly that the pd change is untested and why.
Contributor
Author
There was a problem hiding this comment. Fixed in 9867233. Added a dedicated Paddle adapter test that constructs the backend, calls its eval method, and verifies the normalized frame-major arrays passed to _eval_model. I also corrected the PR body: Paddle is not installed locally, so this test will execute in the Paddle test environment rather than the local consistency loop. Coding agent: Codex |
||
| # by source/tests/common/test_deep_eval_parameter_shorthand.py. | ||
| if backend_name in {"pytorch", "jax", "tf2"} and ( | ||
|
Collaborator
There was a problem hiding this comment.
The consequence is worth deciding deliberately, because #5853 fixes exactly that file with its own inline copy of this ladder. The two PRs touch disjoint files and will merge cleanly, but the result is a shared Cleanest resolutions are to wire dpmodel into the helper here and close #5853, or rebase #5853 onto the helper. Either way, please make the allowlist say why each backend is out, so the gap is a recorded decision rather than an invisible one.
Contributor
Author
There was a problem hiding this comment. Fixed in 9867233. dpmodel now uses the shared helper before auto batching, and the consistency allowlist now exercises every backend present in the loop. I also wired the legacy TensorFlow adapter because it had the same normalization-after-batching ordering. Coding agent: Codex |
||
| 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 +246,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())) | ||
|
Collaborator
There was a problem hiding this comment. Every value in this test is That matters because frame-major and atom-major correctness under batching is the whole point of the change. The sibling test in I raised the same point on #5853 and CodeRabbit flagged this exact block on this PR; it is still unaddressed at head, so flagging once more.
Contributor
Author
There was a problem hiding this comment. Fixed in 9867233. The consistency regression now uses distinct coordinates, boxes, per-atom values, and distinct full frame-major parameter rows. It checks shared-per-atom, shared-all-atoms, and full-frame-major cases, so axis swaps and frame reuse change the numerical result. Coding agent: Codex |
||
| 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: | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
deepmd/pt_expt/infer/deep_eval.pyhas the same defect and is not wired to this helper: rawfparam/aparamgo into_eval_func/AutoBatchSize, and_eval_model/_eval_model_lowerdo barefparam.reshape(nframes, dim_fparam)andaparam.reshape(nframes, natoms, dim_aparam)with no shorthand handling. It is also absent from the io test loop, so nothing would catch it.Its own
evaldocstring documents only the fullnframes x ...forms, which quietly narrows the baseDeepEvalBackend.evalcontract - that has listed all three shorthand forms since #3213 and is the contract this PR is enforcing everywhere else.Given pt_expt is the actively developed backend and carries no back-compat constraints, adding the same one-line call here seems worth doing in this PR rather than leaving a fifth adapter behind.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in 9867233. pt_expt now calls the shared helper before auto batching, its eval docstring documents all supported shorthand forms, and a direct adapter regression verifies the normalized arguments reaching _eval_model.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh