fix(pt_expt): dispatch property models to DeepProperty inference - #5724
Conversation
A pt_expt property checkpoint could be built and trained but not evaluated:
pt_expt DeepEval.model_type dispatched only energy/DOS/dipole/polar/WFC and
raised RuntimeError("Unknown model type") for property outputs, and the
evaluator did not expose the property metadata getters DeepProperty needs.
Underneath, the pt_expt PropertyModel itself only implemented get_var_name, not
get_task_dim or get_intensive.
Add get_task_dim and get_intensive to the pt_expt PropertyModel (mirroring the
PyTorch model), import DeepProperty in the pt_expt evaluator and 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 and getters are guarded by hasattr so non-property and
unknown model types are unaffected (they keep matching their own branches or
fall through to "Unknown model type"), and the guard on self._dpmodel keeps
metadata-only mode (no reconstructed model) raising a clear NotImplementedError
rather than mis-dispatching.
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 the DeepEval
construction raises "Unknown model type". Verified pt property inference works
end to end (same mechanism) and the pt_expt energy inference suite is unaffected.
Fix deepmodeling#5671
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughProperty-model support is added across inference backends and export paths. ChangesProperty model inference support
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt_expt/infer/deep_eval.py (1)
711-718: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueProperty dispatch requires the reconstructed dpmodel; metadata-only
.pte/.pt2property archives still hitRuntimeError("Unknown model type").The new
elifbranch only fires whenself._dpmodel is not None. In metadata-only mode (_init_from_metadata, used whenmodel.jsonis absent), a property model's output type would still be present inmodel_output_type, but there is noget_var_name()to check against, so it falls through toraise RuntimeError("Unknown model type")— the exact error this PR intends to fix. This is consistent with the existing pattern thatget_var_name/get_task_dim/get_intensiveare also unavailable in metadata-only mode, so it's a known limitation rather than a regression, but the resulting generic error message will confuse users who hit it via a metadata-only archive.💡 Optional: give metadata-only property models a clearer error
elif ( self._dpmodel is not None and hasattr(self._dpmodel, "get_var_name") and self._dpmodel.get_var_name() in model_output_type ): return DeepProperty + elif self._dpmodel is None and "task_dim" in self.metadata: + raise NotImplementedError( + "Property model dispatch requires the reconstructed dpmodel " + "(model.json); this archive was loaded in metadata-only mode." + ) else: raise RuntimeError("Unknown model type")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt_expt/infer/deep_eval.py` around lines 711 - 718, The property-type dispatch in the model-type selection logic is too strict because it only checks self._dpmodel.get_var_name() when self._dpmodel is available, so metadata-only .pte/.pt2 archives fall through to the generic RuntimeError("Unknown model type"). Update the dispatch around the self._dpmodel branch to handle metadata-only property archives explicitly, either by recognizing the property model from the available metadata/model_output_type path or by raising a clearer, targeted error from this selector in DeepEval. Keep the fix localized to the model-type resolution logic and preserve the existing behavior for non-property models.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 711-718: The property-type dispatch in the model-type selection
logic is too strict because it only checks self._dpmodel.get_var_name() when
self._dpmodel is available, so metadata-only .pte/.pt2 archives fall through to
the generic RuntimeError("Unknown model type"). Update the dispatch around the
self._dpmodel branch to handle metadata-only property archives explicitly,
either by recognizing the property model from the available
metadata/model_output_type path or by raising a clearer, targeted error from
this selector in DeepEval. Keep the fix localized to the model-type resolution
logic and preserve the existing behavior for non-property models.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5da33ea0-3592-4528-83f9-9761860d1752
📒 Files selected for processing (3)
deepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/property_model.pysource/tests/pt_expt/infer/test_deep_eval_property.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5724 +/- ##
==========================================
- Coverage 79.62% 79.51% -0.12%
==========================================
Files 1014 1014
Lines 115533 115676 +143
Branches 4276 4272 -4
==========================================
- Hits 91995 91981 -14
- Misses 21994 22148 +154
- Partials 1544 1547 +3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Thanks for the fix. I think this should be handled consistently across the DeepEval backends before merging.
This PR adds the DeepProperty dispatch and the property metadata getters for pt_expt, but the same gap still exists in the other backend DeepEval.model_type implementations I checked:
deepmd/dpmodel/infer/deep_eval.pydeepmd/jax/infer/deep_eval.pydeepmd/tf2/infer/deep_eval.py
Those backends still only dispatch energy / dos / dipole / polar / wfc, so a property model reaching the generic DeepEval(...) path would still raise Unknown model type. They also do not currently expose the get_var_name() / get_task_dim() / get_intensive() hooks needed by DeepProperty.change_output_def().
Please extend the same DeepProperty dispatch logic and metadata getter support to dpmodel, jax, and tf2 as well, with tests where the backend supports property inference. If any of these backends intentionally cannot support property models, please make that explicit in the code/tests/docs rather than leaving the same Unknown model type failure mode.
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
PR deepmodeling#5724 added property-model dispatch to the pt_expt evaluator only. Review feedback asked for the same across the remaining backends, whose generic DeepEval path still raised RuntimeError("Unknown model type") for property models and did not expose get_var_name/get_task_dim/get_intensive. Rather than copying the dispatch and getters into each backend, centralize them in the shared DeepEvalBackend base: model_type is now a concrete property that dispatches on the model output types (and the property variable name via get_model()), and get_var_name/get_task_dim/get_intensive delegate to get_model(). dpmodel, jax and tf2 drop their duplicated model_type overrides and inherit this; pt and pt_expt keep their overrides (extra branches / idiosyncratic model access). The property getters live once on the dpmodel PropertyModel base, which the jax and tf2 model classes already subclass, so they inherit them for free. The jax HLO and tf2 SavedModel evaluators wrap an exported artifact with no live model, so the property variable name, output dimension and intensiveness are persisted into the StableHLO constants / SavedModel and read back by the wrappers; HLO.model_output_def rebuilds the dynamic property output def from that metadata. Non-property models store None and are unaffected. Enabling the dispatch surfaced two latent bugs, now fixed: - DeepProperty.eval read the atomic output unconditionally; the dpmodel and jax backends omit it at atomic=False, so it now reads the reduced output for the global property and the atomic output only when atomic=True, mirroring DeepPot.eval. - AutoBatchSize.execute_all unwraps a single-output result out of its tuple, which made the dpmodel/jax/tf2 eval zip iterate over the frame axis; the eval methods now re-wrap it. This also fixes multi-frame global-only DOS inference on these backends, which single-frame tests had masked. Adds full serialize -> artifact -> DeepEval round-trip tests for each backend (dpmodel, jax, tf2), asserting the DeepProperty dispatch, the three getters, and that the multi-frame global property equals the atomic sum.
|
Addressed in 8198ae3 — extended property inference to all three backends (dpmodel, jax, tf2), each with a full serialize → artifact → DeepEval round-trip test asserting the DeepProperty dispatch, the three getters, and that the multi-frame global property equals the atomic sum. Rather than copy the dispatch and getters into each backend, I centralized them in the shared Enabling the dispatch surfaced two latent bugs that single-frame tests had masked, also fixed here: |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/jax/model/hlo.py (1)
248-262: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: align
get_intensivefallback with the other property getters.
get_var_nameandget_task_dimraiseNotImplementedErrorwhen metadata is absent, butget_intensivealways returnsself._intensive(defaulting toFalse), so a non-property HLO model silently reportsintensive=Falseinstead of signaling "not a property model". Since the baseDeepEvalBackend.get_intensiveonly probeshasattr(model, "get_intensive"), this masks the non-property case. Low impact today because it's only exercised for property models, but raising keeps the trio consistent.♻️ Optional consistency tweak
def get_intensive(self) -> bool: """Whether the property is intensive (property models only).""" + if self._var_name is None: + raise NotImplementedError return self._intensive🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/jax/model/hlo.py` around lines 248 - 262, `get_intensive` in the HLO model should match the other property metadata getters by signaling when the model is not a property model. Update the `get_intensive` method in `HLOModel` so it raises `NotImplementedError` when the intensive flag metadata is absent, instead of always returning `self._intensive`; keep the behavior consistent with `get_var_name` and `get_task_dim`, and ensure any property-model initialization path still sets the flag before this getter is used.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@deepmd/jax/model/hlo.py`:
- Around line 248-262: `get_intensive` in the HLO model should match the other
property metadata getters by signaling when the model is not a property model.
Update the `get_intensive` method in `HLOModel` so it raises
`NotImplementedError` when the intensive flag metadata is absent, instead of
always returning `self._intensive`; keep the behavior consistent with
`get_var_name` and `get_task_dim`, and ensure any property-model initialization
path still sets the flag before this getter is used.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 4d153bcd-86af-4799-9895-c980e30c46fb
📒 Files selected for processing (12)
deepmd/dpmodel/infer/deep_eval.pydeepmd/dpmodel/model/property_model.pydeepmd/infer/deep_eval.pydeepmd/infer/deep_property.pydeepmd/jax/infer/deep_eval.pydeepmd/jax/model/hlo.pydeepmd/jax/utils/serialization.pydeepmd/tf2/infer/deep_eval.pydeepmd/tf2/utils/serialization.pysource/tests/common/dpmodel/test_deep_eval_property.pysource/tests/consistent/test_deep_eval_property_tf2.pysource/tests/jax/test_deep_eval_property.py
The JAX HLO and TF2 SavedModel wrappers had byte-identical get_var_name/get_task_dim/get_intensive accessors (return the persisted value, raise if absent), and HLO inlined the property OutputVariableDef construction. Extract a PropertyMetadataHolder mixin in dpmodel/output_def that both wrap in, so the read-back accessors and the property-def builder live once. Each artifact still sets _var_name/_task_dim/_intensive at construction from its own source (StableHLO constants / SavedModel tensors), which is the only genuinely per-backend part; the dpmodel PropertyModel is unchanged since it computes these from the live fitting net rather than from persisted values.
This reverts commit 9a3cc93. The mixin abstracted three one-line accessors that just return a stored value, so there is no logic that can drift out of sync between the JAX HLO and TF2 SavedModel wrappers -- the divergence risk that motivates DRY is essentially zero here. In exchange it added a mixin in dpmodel/output_def (a module about output definitions, not artifact metadata) reached across packages by both wrappers. The two small local getters are self-documenting where they sit and cheaper to read than the shared indirection, so keep them local. The DeepEvalBackend.model_type centralization is retained: that one shares a dispatch decision that must stay consistent across backends.
The test computed INSTALLED_TF2 from a raw Backend availability check, which is True in the main Python CI job (TensorFlow is installed). It therefore ran there and failed with "tf2 backend requires TensorFlow eager execution" — the tf v1 backend disables eager earlier in that job. Import INSTALLED_TF2 from consistent/common instead, which also requires RUN_TF2_BACKEND_TESTS (DP_TEST_TF2_ONLY), so the test only runs in the tf2-only job and is skipped elsewhere.
The three ad-hoc per-backend property tests (dpmodel/jax/tf2) were not consistency tests -- the tf2 one in particular only lived under consistent/ because the tf2 CI job collects that directory. Replace them with a single TestDeepProperty in consistent/io/test_io.py, the genuine cross-backend DeepEval round trip: it serializes one property model, evaluates it through every available backend and asserts they agree to 1e-12. To support this, make IOTest.test_deep_eval arity-agnostic (split global vs atomic outputs by the non-atomic return length instead of a hard-coded 3, so it works for property's single output as well as energy's three) and add a skip_backends hook (tf v1 has no property model). Folding property into test_io surfaced that the jax2tf .savedmodel path (TFModelWrapper) did not carry property metadata, unlike the .hlo path, so it raised "Unknown model type". Persist var_name/task_dim/intensive into the jax2tf SavedModel and read them back in TFModelWrapper, mirroring the .hlo and .savedmodeltf paths, so all jax/tf2 export formats support property inference.
… cycle CodeQL flagged six cyclic imports: the centralized DeepEvalBackend.model_type imported the concrete Deep* wrapper classes (deep_pot, deep_dos, ...), but those modules import DeepEval from deep_eval, so referencing them from the base forms an import cycle regardless of the deferral. The cycle-free home for dispatch that names the Deep* wrappers is each backend's own module, which lives in a separate package the wrappers do not import back. Restore model_type as an abstractmethod on the base and re-add the concrete dispatch (with the property branch) to the dpmodel, jax and tf2 evaluators, importing the wrappers at the top level as before. The shared _get_property_var_name helper and the get_var_name/get_task_dim/get_intensive getters stay on the base -- they delegate to get_model() and import nothing, so they cause no cycle and keep the getter logic in one place.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Thanks, this addresses the earlier cross-backend concern. The property dispatch is now covered for pt_expt, dpmodel, jax, and tf2; the single-output eval wrapping issue is handled; and the cross-backend IO test covers the intended round trip while skipping tf v1 where unsupported.
CI is green on the latest commit, and the prior CodeQL cyclic-import issue is resolved by keeping dispatch in each backend module.
Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
…ty-eval # Conflicts: # deepmd/jax/model/hlo.py
4 participants
Problem
Fixes #5671. A pt_expt property checkpoint could be constructed and trained by the backend stack but not evaluated.
pt_exptDeepEval.model_typedispatched only energy, DOS, dipole, polar, and WFC outputs and then raisedRuntimeError("Unknown model type")for a property model, and the evaluator did not expose theget_intensive/get_var_name/get_task_dimgetters thatDeepPropertyneeds. Underneath, the pt_exptPropertyModelitself only implementedget_var_name, notget_task_dimorget_intensive, so simply mirroring the dispatch would not have been enough.Fix
pt_expt/model/property_model.py: addget_task_dim(fitting output dimension) andget_intensive(from the output def), mirroring the PyTorch property model.pt_expt/infer/deep_eval.py: importDeepProperty, dispatch to it when the property variable name appears in the model output, and exposeget_var_name/get_task_dim/get_intensivethat delegate to the reconstructed model.The dispatch branch and the getters are guarded by
hasattr(and byself._dpmodel is not None), so:"Unknown model type"rather than raising anAttributeError;NotImplementedErrorfor the property getters instead of mis-dispatching.Test
Adds
source/tests/pt_expt/infer/test_deep_eval_property.py, a fullserialize -> .pte -> DeepEvalround trip assertingmodel_type is DeepProperty, the three getters, and the eval output shape. Without the fix, constructing theDeepEvalraisesRuntimeError("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.Summary by CodeRabbit