Skip to content

Commit 1fa7f71

Browse files
wanghan-iapcmHan Wang
andauthored
fix(pt_expt): dispatch property models to DeepProperty inference (deepmodeling#5724)
## Problem Fixes deepmodeling#5671. A pt_expt property checkpoint could be constructed and trained by the backend stack but not evaluated. `pt_expt` `DeepEval.model_type` dispatched only energy, DOS, dipole, polar, and WFC outputs and then raised `RuntimeError("Unknown model type")` for a property model, and the evaluator did not expose the `get_intensive` / `get_var_name` / `get_task_dim` getters that `DeepProperty` needs. Underneath, the pt_expt `PropertyModel` itself only implemented `get_var_name`, not `get_task_dim` or `get_intensive`, so simply mirroring the dispatch would not have been enough. ## Fix - `pt_expt/model/property_model.py`: add `get_task_dim` (fitting output dimension) and `get_intensive` (from the output def), mirroring the PyTorch property model. - `pt_expt/infer/deep_eval.py`: import `DeepProperty`, dispatch to it when the property variable name appears in the model output, and expose `get_var_name` / `get_task_dim` / `get_intensive` that delegate to the reconstructed model. The dispatch branch and the getters are guarded by `hasattr` (and by `self._dpmodel is not None`), so: - energy/DOS/dipole/polar/WFC models keep matching their own branches first; - genuinely unknown model types still fall through to `"Unknown model type"` rather than raising an `AttributeError`; - metadata-only mode (no reconstructed dpmodel, used by the C++ AOTI path) raises a clear `NotImplementedError` for the property getters instead of mis-dispatching. ## Test Adds `source/tests/pt_expt/infer/test_deep_eval_property.py`, a full `serialize -> .pte -> DeepEval` round trip asserting `model_type is DeepProperty`, the three getters, and the eval output shape. Without the fix, constructing the `DeepEval` raises `RuntimeError("Unknown model type")`. Verification for the reviewer's peace of mind: pt-backend property inference works end to end via the same mechanism (dispatch + getters + eval), and the full pt_expt energy inference suite (`source/tests/pt_expt/infer/test_deep_eval.py`, 92 passed / 2 skipped) is unaffected by the added branch. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for property-style models across inference backends, including correct model dispatch to property evaluation. * Exposed property metadata (variable name, task dimension, and intensive/extensive flag) via new public getters and preserved it through export/import. * **Bug Fixes** * Fixed single-output evaluation so results are wrapped consistently and mapped correctly to requested outputs. * Corrected atomic property tensor reshaping behavior for non-atomic evaluation. * **Tests** * Added end-to-end unit, JAX, and TF2 consistency tests covering dispatch, metadata getters, and output shapes/values. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
1 parent f73de32 commit 1fa7f71

15 files changed

Lines changed: 474 additions & 17 deletions

File tree

deepmd/dpmodel/infer/deep_eval.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@
4747
from deepmd.infer.deep_pot import (
4848
DeepPot,
4949
)
50+
from deepmd.infer.deep_property import (
51+
DeepProperty,
52+
)
5053
from deepmd.infer.deep_wfc import (
5154
DeepWFC,
5255
)
@@ -134,8 +137,9 @@ def has_default_fparam(self) -> bool:
134137

135138
@property
136139
def model_type(self) -> type["DeepEvalWrapper"]:
137-
"""The the evaluator of the model type."""
138-
model_output_type = self.dp.model_output_type()
140+
"""The evaluator of the model type."""
141+
model = self.get_model()
142+
model_output_type = model.model_output_type()
139143
if "energy" in model_output_type:
140144
return DeepPot
141145
elif "dos" in model_output_type:
@@ -146,6 +150,8 @@ def model_type(self) -> type["DeepEvalWrapper"]:
146150
return DeepPolar
147151
elif "wfc" in model_output_type:
148152
return DeepWFC
153+
elif self._get_property_var_name(model) in model_output_type:
154+
return DeepProperty
149155
else:
150156
raise RuntimeError("Unknown model type")
151157

@@ -238,6 +244,12 @@ def eval(
238244
out = self._eval_func(self._eval_model, numb_test, natoms)(
239245
coords, cells, atom_types, fparam, aparam, request_defs
240246
)
247+
# ``AutoBatchSize.execute_all`` unwraps a single-output result out of
248+
# its tuple, which would make ``zip`` iterate over the array's frame
249+
# axis. Re-wrap so the request-def names line up (a single request def
250+
# arises for global-only DOS/property inference at atomic=False).
251+
if not isinstance(out, tuple):
252+
out = (out,)
241253
return dict(
242254
zip(
243255
[x.name for x in request_defs],

deepmd/dpmodel/model/property_model.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,14 @@ def get_var_name(self) -> str:
4343
"""Get the name of the property."""
4444
return self.get_fitting_net().var_name
4545

46+
def get_task_dim(self) -> int:
47+
"""Get the output dimension of the property."""
48+
return self.get_fitting_net().dim_out
49+
50+
def get_intensive(self) -> bool:
51+
"""Whether the property is intensive."""
52+
return self.model_output_def()[self.get_var_name()].intensive
53+
4654
def call(
4755
self,
4856
coord: Array,

deepmd/infer/deep_eval.py

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,31 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool:
375375
@property
376376
@abstractmethod
377377
def model_type(self) -> type["DeepEval"]:
378-
"""The the evaluator of the model type."""
378+
"""The evaluator of the model type.
379+
380+
Each backend implements the dispatch on its own module so it can import
381+
the concrete ``Deep*`` wrapper classes at the top level. Those wrappers
382+
import ``DeepEval`` from this module, so a dispatch here would form an
383+
import cycle (flagged by CodeQL). :meth:`_get_property_var_name` is
384+
provided for the shared property branch.
385+
"""
386+
387+
@staticmethod
388+
def _get_property_var_name(model: Any) -> str | None:
389+
"""Return the property variable name of ``model``, or ``None``.
390+
391+
Used by every backend's ``model_type`` to detect a property model.
392+
``get_var_name`` may be absent (dpmodel/pt live models expose it only on
393+
property models) or present-but-unimplemented (jax/tf2 artifacts always
394+
define it and raise ``NotImplementedError`` otherwise), so probe
395+
defensively.
396+
"""
397+
if not hasattr(model, "get_var_name"):
398+
return None
399+
try:
400+
return model.get_var_name()
401+
except NotImplementedError:
402+
return None
379403

380404
@abstractmethod
381405
def get_sel_type(self) -> list[int]:
@@ -414,7 +438,24 @@ def get_has_hessian(self) -> bool:
414438
return False
415439

416440
def get_var_name(self) -> str:
417-
"""Get the name of the fitting property."""
441+
"""Get the name of the fitting property (property models only)."""
442+
model = self.get_model()
443+
if hasattr(model, "get_var_name"):
444+
return model.get_var_name()
445+
raise NotImplementedError
446+
447+
def get_task_dim(self) -> int:
448+
"""Get the output dimension of the property (property models only)."""
449+
model = self.get_model()
450+
if hasattr(model, "get_task_dim"):
451+
return model.get_task_dim()
452+
raise NotImplementedError
453+
454+
def get_intensive(self) -> bool:
455+
"""Whether the property is intensive (property models only)."""
456+
model = self.get_model()
457+
if hasattr(model, "get_intensive"):
458+
return model.get_intensive()
418459
raise NotImplementedError
419460

420461
@abstractmethod

deepmd/infer/deep_property.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -134,14 +134,13 @@ def eval(
134134
aparam=aparam,
135135
**kwargs,
136136
)
137-
atomic_property = results[self.get_var_name()].reshape(
138-
nframes, natoms, self.get_task_dim()
139-
)
140137
property = results[f"{self.get_var_name()}_redu"].reshape(
141138
nframes, self.get_task_dim()
142139
)
143-
144140
if atomic:
141+
atomic_property = results[self.get_var_name()].reshape(
142+
nframes, natoms, self.get_task_dim()
143+
)
145144
return (
146145
property,
147146
atomic_property,

deepmd/jax/infer/deep_eval.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@
4141
from deepmd.infer.deep_pot import (
4242
DeepPot,
4343
)
44+
from deepmd.infer.deep_property import (
45+
DeepProperty,
46+
)
4447
from deepmd.infer.deep_wfc import (
4548
DeepWFC,
4649
)
@@ -152,7 +155,8 @@ def get_dim_aparam(self) -> int:
152155
@property
153156
def model_type(self) -> type["DeepEvalWrapper"]:
154157
"""The evaluator of the model type."""
155-
model_output_type = self.dp.model_output_type()
158+
model = self.get_model()
159+
model_output_type = model.model_output_type()
156160
if "energy" in model_output_type:
157161
return DeepPot
158162
elif "dos" in model_output_type:
@@ -163,6 +167,8 @@ def model_type(self) -> type["DeepEvalWrapper"]:
163167
return DeepPolar
164168
elif "wfc" in model_output_type:
165169
return DeepWFC
170+
elif self._get_property_var_name(model) in model_output_type:
171+
return DeepProperty
166172
else:
167173
raise RuntimeError("Unknown model type")
168174

@@ -270,6 +276,12 @@ def eval(
270276
out = self._eval_func(self._eval_model, numb_test, natoms)(
271277
coords, cells, atom_types, fparam, aparam, request_defs
272278
)
279+
# ``AutoBatchSize.execute_all`` unwraps a single-output result out of
280+
# its tuple, which would make ``zip`` iterate over the array's frame
281+
# axis. Re-wrap so the request-def names line up (a single request def
282+
# arises for global-only DOS/property inference at atomic=False).
283+
if not isinstance(out, tuple):
284+
out = (out,)
273285
return dict(
274286
zip(
275287
[x.name for x in request_defs],

deepmd/jax/jax2tf/serialization.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,28 @@ def get_default_fparam() -> tf.Tensor:
318318

319319
tf_model.get_default_fparam = get_default_fparam
320320

321+
# property models: persist the output name/dimension/intensiveness so
322+
# the evaluator can dispatch to DeepProperty and reshape the output.
323+
if hasattr(model, "get_var_name"):
324+
325+
@tf.function
326+
def get_var_name() -> tf.Tensor:
327+
return tf.constant(model.get_var_name(), dtype=tf.string)
328+
329+
tf_model.get_var_name = get_var_name
330+
331+
@tf.function
332+
def get_task_dim() -> tf.Tensor:
333+
return tf.constant(model.get_task_dim(), dtype=tf.int64)
334+
335+
tf_model.get_task_dim = get_task_dim
336+
337+
@tf.function
338+
def get_intensive() -> tf.Tensor:
339+
return tf.constant(model.get_intensive(), dtype=tf.bool)
340+
341+
tf_model.get_intensive = get_intensive
342+
321343
tf.saved_model.save(
322344
tf_model,
323345
model_file,

deepmd/jax/jax2tf/tfmodel.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,15 @@ def __init__(
7878
self.default_fparam = self.model.get_default_fparam().numpy().tolist()
7979
else:
8080
self.default_fparam = None
81+
# property models only (absent for other model types).
82+
if hasattr(self.model, "get_var_name"):
83+
self._var_name = self.model.get_var_name().numpy().decode()
84+
self._task_dim = self.model.get_task_dim().numpy().item()
85+
self._intensive = self.model.get_intensive().numpy().item()
86+
else:
87+
self._var_name = None
88+
self._task_dim = None
89+
self._intensive = False
8190

8291
def __call__(
8392
self,
@@ -175,9 +184,27 @@ def call(
175184

176185
def model_output_def(self) -> ModelOutputDef:
177186
return ModelOutputDef(
178-
FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()])
187+
FittingOutputDef(
188+
[self._output_var_def(tt) for tt in self.model_output_type()]
189+
)
179190
)
180191

192+
def _output_var_def(self, name: str) -> OutputVariableDef:
193+
if name in OUTPUT_DEFS:
194+
return OUTPUT_DEFS[name]
195+
# property models carry a user-defined output name (``var_name``) that
196+
# is not in the fixed table; rebuild its def from the persisted metadata.
197+
if self._var_name is not None and name == self._var_name:
198+
return OutputVariableDef(
199+
self._var_name,
200+
shape=[self._task_dim],
201+
reducible=True,
202+
r_differentiable=False,
203+
c_differentiable=False,
204+
intensive=self._intensive,
205+
)
206+
raise KeyError(f"Unknown model output variable {name!r}")
207+
181208
def call_lower(
182209
self,
183210
extended_coord: jnp.ndarray,
@@ -349,3 +376,19 @@ def has_default_fparam(self) -> bool:
349376
def get_default_fparam(self) -> list[float] | None:
350377
"""Get the default frame parameters."""
351378
return self.default_fparam
379+
380+
def get_var_name(self) -> str:
381+
"""Get the name of the property (property models only)."""
382+
if self._var_name is None:
383+
raise NotImplementedError
384+
return self._var_name
385+
386+
def get_task_dim(self) -> int:
387+
"""Get the output dimension of the property (property models only)."""
388+
if self._task_dim is None:
389+
raise NotImplementedError
390+
return self._task_dim
391+
392+
def get_intensive(self) -> bool:
393+
"""Whether the property is intensive (property models only)."""
394+
return self._intensive

deepmd/jax/model/hlo.py

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ def __init__(
6969
has_default_fparam: bool = False,
7070
default_fparam: list[float] | None = None,
7171
numb_dos: int = 0,
72+
# property models only
73+
var_name: str | None = None,
74+
task_dim: int | None = None,
75+
intensive: bool = False,
7276
) -> None:
7377
self._call_lower = jax_export.deserialize(stablehlo).call
7478
self._call_lower_atomic_virial = jax_export.deserialize(
@@ -93,6 +97,9 @@ def __init__(
9397
self._has_default_fparam = has_default_fparam
9498
self.default_fparam = default_fparam
9599
self.numb_dos = numb_dos
100+
self._var_name = var_name
101+
self._task_dim = task_dim
102+
self._intensive = intensive
96103

97104
def __call__(
98105
self,
@@ -180,9 +187,27 @@ def call(
180187

181188
def model_output_def(self) -> ModelOutputDef:
182189
return ModelOutputDef(
183-
FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()])
190+
FittingOutputDef(
191+
[self._output_var_def(tt) for tt in self.model_output_type()]
192+
)
184193
)
185194

195+
def _output_var_def(self, name: str) -> OutputVariableDef:
196+
if name in OUTPUT_DEFS:
197+
return OUTPUT_DEFS[name]
198+
# property models carry a user-defined output name (``var_name``) that
199+
# is not in the fixed table; rebuild its def from the persisted metadata.
200+
if self._var_name is not None and name == self._var_name:
201+
return OutputVariableDef(
202+
self._var_name,
203+
shape=[self._task_dim],
204+
reducible=True,
205+
r_differentiable=False,
206+
c_differentiable=False,
207+
intensive=self._intensive,
208+
)
209+
raise KeyError(f"Unknown model output variable {name!r}")
210+
186211
def call_lower(
187212
self,
188213
extended_coord: jnp.ndarray,
@@ -233,6 +258,22 @@ def get_dim_aparam(self) -> int:
233258
"""Get the number (dimension) of atomic parameters of this atomic model."""
234259
return self.dim_aparam
235260

261+
def get_var_name(self) -> str:
262+
"""Get the name of the property (property models only)."""
263+
if self._var_name is None:
264+
raise NotImplementedError
265+
return self._var_name
266+
267+
def get_task_dim(self) -> int:
268+
"""Get the output dimension of the property (property models only)."""
269+
if self._task_dim is None:
270+
raise NotImplementedError
271+
return self._task_dim
272+
273+
def get_intensive(self) -> bool:
274+
"""Whether the property is intensive (property models only)."""
275+
return self._intensive
276+
236277
def get_sel_type(self) -> list[int]:
237278
"""Get the selected atom types of this model.
238279

deepmd/jax/utils/serialization.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,18 @@ def call_lower_with_fixed_do_atomic_virial(
321321
"sel": model.get_sel(),
322322
"has_default_fparam": model.has_default_fparam(),
323323
"default_fparam": model.get_default_fparam(),
324+
# property models: the output name/dimension/intensiveness cannot be
325+
# recovered from the StableHLO alone, so persist them for the
326+
# evaluator (None for non-property models).
327+
"var_name": model.get_var_name()
328+
if hasattr(model, "get_var_name")
329+
else None,
330+
"task_dim": model.get_task_dim()
331+
if hasattr(model, "get_task_dim")
332+
else None,
333+
"intensive": model.get_intensive()
334+
if hasattr(model, "get_intensive")
335+
else False,
324336
}
325337
save_dp_model(filename=model_file, model_dict=data)
326338
elif model_file.endswith(".savedmodel"):

0 commit comments

Comments
 (0)