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,75 @@ | |
| 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. | ||
|
|
||
| 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) | ||
| 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 +1017,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 |
|---|---|---|
|
|
@@ -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 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 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( | ||
|
Collaborator
There was a problem hiding this comment. This wires The summary claims the change covers "normal, spin, embedding, descriptor, and fitting-last-layer routes". For pt_expt the descriptor and fitting-last-layer routes are not covered: if aparam is not None:
aparam_t = torch.tensor(
aparam.reshape(nframes, natoms, self.get_dim_aparam()),So a backend-direct I want to be clear this is milder than the bug the PR fixes: these routes do not go through Normalizing inside On coverage while you are here: the new pd and pt_expt adapter tests both set
Contributor
Author
There was a problem hiding this comment. Fixed in 866ce34. Coding agent: opencode |
||
| 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: | ||
|
|
@@ -2203,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, | ||
|
|
@@ -2274,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, | ||
|
|
||
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