feat(dpa3): decouple charge_spin from fparam#5431
Open
iProzd wants to merge 1 commit intodeepmodeling:masterfrom
Open
feat(dpa3): decouple charge_spin from fparam#5431iProzd wants to merge 1 commit intodeepmodeling:masterfrom
iProzd wants to merge 1 commit intodeepmodeling:masterfrom
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Comment on lines
+40
to
+48
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Comment on lines
+70
to
+79
| return model_predict | ||
|
|
||
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, |
Comment on lines
+40
to
+48
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Comment on lines
+64
to
+73
| return model_predict | ||
|
|
||
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, |
Comment on lines
+40
to
+48
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Comment on lines
+70
to
+79
| return model_predict | ||
|
|
||
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, |
Comment on lines
+40
to
+48
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Comment on lines
+64
to
+73
| return model_predict | ||
|
|
||
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, |
Comment on lines
+44
to
+52
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Comment on lines
+69
to
+78
| return model_predict | ||
|
|
||
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR makes charge/spin a first-class per-frame input for DPA3 by introducing a dedicated charge_spin kwarg through the forward/eval/training chains, removing the prior coupling to fparam, and adding an optional default_chg_spin fallback on the DPA3 descriptor.
Changes:
- Plumb
charge_spin: Tensor | Nonethrough PT / dpmodel / PT_EXPT forward paths, plus data loading/statistics and CLI evaluation entrypoints. - Add
default_chg_spinto DPA3 descriptor config/serialization and wire a default fallback indp_atomic_model.forward_atomic. - Update/expand unit and consistency tests to cover no/explicit/default charge+spin modes.
Reviewed changes
Copilot reviewed 60 out of 60 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| source/tests/universal/dpmodel/model/test_model.py | Update universal skip logic to require default_chg_spin when add_chg_spin_ebd is enabled. |
| source/tests/universal/dpmodel/descriptor/test_descriptor.py | Add default_chg_spin parameterization for DPA3 universal descriptor tests. |
| source/tests/pt/model/test_dpa3.py | Rework PT DPA3 consistency test across three cs_mode cases and switch inputs from fparam to charge_spin. |
| source/tests/pt_expt/descriptor/test_dpa3.py | Add PT_EXPT DPA3 charge/spin consistency test vs dpmodel, including default mode. |
| source/tests/consistent/model/test_ener.py | Reparameterize energy consistency test across charge/spin modes and switch plumbing to charge_spin. |
| source/tests/consistent/descriptor/test_dpa3.py | Switch consistent descriptor tests from fparam to charge_spin. |
| source/tests/consistent/descriptor/common.py | Thread charge_spin through backend-agnostic descriptor evaluation helpers. |
| deepmd/utils/argcheck.py | Add default_chg_spin argument and update add_chg_spin_ebd documentation for DPA3. |
| deepmd/pt/utils/stat.py | Pass charge_spin through PT stat prediction path. |
| deepmd/pt/train/wrapper.py | Accept and forward charge_spin through the PT training wrapper. |
| deepmd/pt/train/training.py | Add charge_spin input key handling and register it as a data requirement when needed. |
| deepmd/pt/model/model/make_model.py | Forward charge_spin through common forward paths and expose new chg/spin-related query methods. |
| deepmd/pt/model/model/ener_model.py | Add charge_spin to PT energy model forward/forward_lower signatures and calls. |
| deepmd/pt/model/descriptor/se_t.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/descriptor/se_t_tebd.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/descriptor/se_r.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/descriptor/se_a.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/descriptor/hybrid.py | Forward charge_spin into sub-descriptors during hybrid descriptor evaluation. |
| deepmd/pt/model/descriptor/dpa3.py | Implement default_chg_spin config/serialization + consume charge_spin instead of fparam. |
| deepmd/pt/model/descriptor/dpa2.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/descriptor/dpa1.py | Add charge_spin passthrough to descriptor forward signature. |
| deepmd/pt/model/atomic_model/pairtab_atomic_model.py | Add charge_spin passthrough to atomic model forward signature. |
| deepmd/pt/model/atomic_model/linear_atomic_model.py | Add charge_spin passthrough and forward it into underlying atomic model. |
| deepmd/pt/model/atomic_model/dp_atomic_model.py | Apply default_chg_spin fallback and pass charge_spin into the descriptor. |
| deepmd/pt/model/atomic_model/base_atomic_model.py | Add base hooks and forward plumbing for charge_spin across atomic-model interfaces. |
| deepmd/pt/infer/deep_eval.py | Add charge_spin to PT deep_eval eval path and model call plumbing. |
| deepmd/pt_expt/train/wrapper.py | Accept and forward charge_spin through the PT_EXPT training wrapper. |
| deepmd/pt_expt/train/training.py | Add charge/spin data requirement and thread charge_spin through trace/compile + runtime get_data. |
| deepmd/pt_expt/model/spin_ener_model.py | Add charge_spin to PT_EXPT spin-energy model forward chains and exportable tracing fn signatures. |
| deepmd/pt_expt/model/property_model.py | Add charge_spin to PT_EXPT property model forward/lower/exportable paths. |
| deepmd/pt_expt/model/polar_model.py | Add charge_spin to PT_EXPT polar model forward/lower/exportable paths. |
| deepmd/pt_expt/model/make_model.py | Thread charge_spin through PT_EXPT common/atomic forward and Hessian helpers. |
| deepmd/pt_expt/model/ener_model.py | Add charge_spin to PT_EXPT energy model forward/lower/exportable paths. |
| deepmd/pt_expt/model/dp_zbl_model.py | Add charge_spin to PT_EXPT ZBL model forward/lower/exportable paths. |
| deepmd/pt_expt/model/dp_linear_model.py | Add charge_spin to PT_EXPT linear model forward/lower/exportable paths. |
| deepmd/pt_expt/model/dos_model.py | Add charge_spin to PT_EXPT DOS model forward/lower/exportable paths. |
| deepmd/pt_expt/model/dipole_model.py | Add charge_spin to PT_EXPT dipole model forward/lower/exportable paths. |
| deepmd/pt_expt/infer/deep_eval.py | Add charge_spin to PT_EXPT evaluator API, metadata, and input preparation (incl. default fallback). |
| deepmd/infer/deep_eval.py | Add chg/spin capability query hooks on the backend-agnostic evaluator interface. |
| deepmd/entrypoints/test.py | Extend dp test to request/pass charge_spin when required by the model. |
| deepmd/dpmodel/utils/stat.py | Pass charge_spin through dpmodel stat prediction path. |
| deepmd/dpmodel/utils/lmdb_data.py | Ensure LMDB frames always get find_charge_spin alongside other optional inputs. |
| deepmd/dpmodel/utils/batch.py | Treat charge_spin as a model input key in batch normalization/splitting. |
| deepmd/dpmodel/model/make_model.py | Thread charge_spin through dpmodel forward/call_common interfaces and add chg/spin query methods. |
| deepmd/dpmodel/model/ener_model.py | Add charge_spin to dpmodel energy model call/call_lower signatures and calls. |
| deepmd/dpmodel/descriptor/se_t.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/descriptor/se_t_tebd.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/descriptor/se_r.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/descriptor/se_e2_a.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/descriptor/make_base_descriptor.py | Update base descriptor interface to accept charge_spin. |
| deepmd/dpmodel/descriptor/hybrid.py | Forward charge_spin into sub-descriptors during hybrid descriptor evaluation. |
| deepmd/dpmodel/descriptor/dpa3.py | Implement default_chg_spin config/serialization + consume charge_spin instead of fparam. |
| deepmd/dpmodel/descriptor/dpa2.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/descriptor/dpa1.py | Add charge_spin passthrough to dpmodel descriptor call signature. |
| deepmd/dpmodel/atomic_model/pairtab_atomic_model.py | Add charge_spin passthrough to dpmodel atomic model forward signature. |
| deepmd/dpmodel/atomic_model/make_base_atomic_model.py | Update base atomic model interface to accept charge_spin. |
| deepmd/dpmodel/atomic_model/linear_atomic_model.py | Add charge_spin passthrough and forward it into underlying atomic model. |
| deepmd/dpmodel/atomic_model/dp_atomic_model.py | Apply default_chg_spin fallback and pass charge_spin into the descriptor; add chg/spin query methods. |
| deepmd/dpmodel/atomic_model/base_atomic_model.py | Add base hooks and forward plumbing for charge_spin across dpmodel atomic-model interfaces. |
| deepmd/calculator.py | Allow ASE calculator path to pass atoms.info["charge_spin"] into evaluator. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+129
to
+131
| assert len(default_chg_spin) == 2, ( | ||
| "default_chg_spin must be a list of length 2 [charge, spin]." | ||
| ) |
Comment on lines
+273
to
+276
| return torch.tensor( | ||
| self.default_chg_spin, | ||
| dtype=self.prec, | ||
| device=env.DEVICE, |
| self.use_econf_tebd = use_econf_tebd | ||
| self.add_chg_spin_ebd = add_chg_spin_ebd | ||
| self.default_chg_spin = default_chg_spin | ||
| if self.add_chg_spin_ebd and self.default_chg_spin is not None: |
Comment on lines
+558
to
+571
| @torch.jit.export | ||
| def has_chg_spin_ebd(self) -> bool: | ||
| """Check if the model has charge spin embedding.""" | ||
| return self.atomic_model.has_chg_spin_ebd() | ||
|
|
||
| @torch.jit.export | ||
| def has_default_chg_spin(self) -> bool: | ||
| """Check if the model has default charge_spin values.""" | ||
| return self.atomic_model.has_default_chg_spin() | ||
|
|
||
| @torch.jit.export | ||
| def get_default_chg_spin(self) -> torch.Tensor | None: | ||
| """Get the default charge_spin values.""" | ||
| return self.atomic_model.get_default_chg_spin() |
Comment on lines
+711
to
+721
| def has_chg_spin_ebd(self) -> bool: | ||
| """Check if the model has charge spin embedding.""" | ||
| return self.atomic_model.has_chg_spin_ebd() | ||
|
|
||
| def has_default_chg_spin(self) -> bool: | ||
| """Check if the model has default charge_spin values.""" | ||
| return self.atomic_model.has_default_chg_spin() | ||
|
|
||
| def get_default_chg_spin(self) -> list[float] | None: | ||
| """Get the default charge_spin values.""" | ||
| return self.atomic_model.get_default_chg_spin() |
Contributor
📝 WalkthroughWalkthroughThis PR adds dedicated ChangesCharge-Spin Input Threading
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~35 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
deepmd/pt/model/descriptor/dpa3.py (1)
206-211:⚠️ Potential issue | 🟠 Major | ⚡ Quick winOff-by-one in charge embedding capacity — charge=100 causes
IndexErrorat runtime.
TypeEmbedNet(200, ...)creates an embedding with valid indices 0–199. After the+ 100offset, the maximum valid input charge value is 99 (index 199), not 100. Passingcharge=100produces index 200, which is out-of-bounds and raises a runtimeIndexError. Yet the comment directly above claims"-100 ~ 100 is a conservative bound", so the intent was clearly to support 100 as a valid value.Fix: use 201 entries to honour the documented inclusive range
[-100, 100]:🐛 Proposed fix
- # -100 ~ 100 is a conservative bound - self.chg_embedding = TypeEmbedNet( - 200, + # -100 ~ 100 inclusive, shifted by +100 → indices 0..200 + self.chg_embedding = TypeEmbedNet( + 201, self.tebd_dim, ... )And correspondingly in the forward:
- charge = charge_spin[:, 0].to(dtype=torch.int64) + 100 + charge = (charge_spin[:, 0].to(dtype=torch.int64) + 100).clamp(0, 200)The same concern applies to the spin embedding (
TypeEmbedNet(100, ...), valid indices 0–99): there is no guard against negative spin values or spin ≥ 100, and no documentation of the valid range. Consider adding a clamp or assert.Also applies to: 586-597
🤖 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/model/descriptor/dpa3.py` around lines 206 - 211, The charge embedding chg_embedding currently constructs TypeEmbedNet(200, ...) while inputs are shifted by +100, causing an off-by-one (charge=100 -> index 200) and IndexError; change the embedding size to 201 so indices 0..200 cover charges -100..100 and ensure the forward mapping that adds 100 uses that new range. Also audit the spin embedding (TypeEmbedNet(100, ...) and any uses that index it) and add a clamp or assert in the forward method to enforce valid spin and charge ranges (or expand the embedding size if the intended inclusive range is larger); apply the same fix to the other occurrence around the spin/charge embedding block referenced later (lines ~586-597).deepmd/pt_expt/model/spin_ener_model.py (1)
151-188:⚠️ Potential issue | 🔴 CriticalSpinModel.forward_common_lower_exportable must accept and forward charge_spin parameter
SpinModel.forward_common_lower_exportable(lines 51–126 inspin_model.py) does not includecharge_spinin its signature or pass it to the traced inner function. However,SpinEnerModel.forward_lower_exportablecalls it withcharge_spin=charge_spinand then attempts to call the returnedtracedmodule withcharge_spinas the 8th positional argument.This creates a critical mismatch:
- The call
self.forward_common_lower_exportable(..., charge_spin=charge_spin)passescharge_spinas an unexpected keyword argument (will be absorbed into**make_fx_kwargsand cause an error inmake_fx)- The traced function was created without
charge_spinin its signature, so callingtraced(..., charge_spin)will fail with a positional argument count mismatchUpdate
SpinModel.forward_common_lower_exportableto match the pattern inmake_model.py(lines 333–435): addcharge_spinto the signature, include it in the innerfnparameters as the 7th positional argument (beforedo_atomic_virial), and pass it through the traced call chain.🤖 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/model/spin_ener_model.py` around lines 151 - 188, SpinModel.forward_common_lower_exportable currently omits the charge_spin parameter causing a signature mismatch when SpinEnerModel.forward_lower_exportable passes charge_spin through; update SpinModel.forward_common_lower_exportable to add charge_spin to its signature, include charge_spin as the 7th positional parameter in the inner traced wrapper (the inner fn that currently accepts extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam, ...) and ensure charge_spin is forwarded into the traced(...) call and through any make_fx_kwargs plumbing (so it is not swallowed by **make_fx_kwargs). Reference SpinModel.forward_common_lower_exportable, SpinEnerModel.forward_lower_exportable, the traced object and inner fn to find and make the change.deepmd/pt_expt/infer/deep_eval.py (1)
1117-1141:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winThe positional vs. keyword
charge_spinasymmetry between.pt2and eager paths is correctly identified.Both the non-spin (
_eval_model) and spin (_eval_model_spin) paths show this pattern:_pt2_runnerreceivescharge_spin_tas a positional argument (7th for non-spin, 8th for spin afterext_spin_t), whileexported_modulereceives it as a keyword argument. The export-time signatures inener_model.py(non-spin) andspin_ener_model.py(spin) correctly definecharge_spinin these exact positions, so the current code is safe. However, the concern about fragility is valid—reordering arguments at export time or inserting new ones beforecharge_spinwould silently pass wrong tensors to.pt2runners. Adding an end-to-end test that exercises.pt2with non-Nonecharge_spin` would catch such regressions.🤖 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 1117 - 1141, The positional vs keyword asymmetry can lead to silent mis-wiring of arguments; update calls to _pt2_runner (in both _eval_model and _eval_model_spin) to pass charge_spin explicitly as a keyword (charge_spin=charge_spin_t) instead of as a positional arg (and similarly ensure ext_spin_t remains positional in the spin path), and add an end-to-end test that invokes the .pt2 runner with a non-None charge_spin to catch future argument-reordering regressions.
🧹 Nitpick comments (5)
deepmd/dpmodel/model/make_model.py (1)
301-334: ⚡ Quick win
charge_spinbypasses_input_type_cast, risking dtype mismatchIn
call_common(line 318) andcall_common_lower(line 387–388),fparamandaparamare explicitly cast to the global float precision through_input_type_cast.charge_spinis forwarded uncasted. If a caller suppliescharge_spinwith a different dtype thancoord, the descriptor receives inconsistently typed inputs, leading to runtime dtype errors in strict backends (e.g., PyTorch with mixed float32/float64).🛠️ Minimal inline fix in `call_common_lower`
cc_ext, _, fp, ap, input_prec = self._input_type_cast( extended_coord, fparam=fparam, aparam=aparam ) del extended_coord, fparam, aparam +if charge_spin is not None: + xp_cs = array_api_compat.array_namespace(charge_spin) + global_dtype = get_xp_precision( + xp_cs, RESERVED_PRECISION_DICT[self.global_np_float_precision] + ) + if charge_spin.dtype != global_dtype: + charge_spin = xp_cs.astype(charge_spin, global_dtype) model_predict = self.forward_common_atomic(A symmetric fix should be applied in
call_commonas well (beforemodel_call_from_call_lower).🤖 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/dpmodel/model/make_model.py` around lines 301 - 334, The call path forwards charge_spin without type normalization, so modify call_common and call_common_lower to include charge_spin in the inputs processed by _input_type_cast (or explicitly cast charge_spin to the same float dtype as returned by _input_type_cast) so that charge_spin is converted to input_prec before calling model_call_from_call_lower or passing into the lower-level descriptor; update the call sites: in call_common (before invoking model_call_from_call_lower) and in call_common_lower (at the start of the function) to use the normalized/converted charge_spin variable returned or produced by _input_type_cast, ensuring all arrays (coord/box/fparam/aparam/charge_spin) share the same dtype.deepmd/pt_expt/infer/deep_eval.py (3)
1273-1297: ⚖️ Poor tradeoffCharge_spin handling duplicated between
_eval_model_spinand_prepare_inputs.
_eval_model_spinre-implements its full input prep inline (including this newcharge_spinblock) instead of going through_prepare_inputs. The block is byte-identical to the one at lines 1057-1081, so any future change (e.g., to the default-fallback policy or dtype) has to be made in two places and will silently drift if missed. This mirrors the pre-existingfparam/aparamduplication in the same method, so it's not a regression — but it's worth refactoring_eval_model_spinto share the prep with_prepare_inputs(e.g., by extracting an_extend_spinhelper to handle the spin-specific part and delegating the rest). Deferable; flagging for tracking.🤖 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 1273 - 1297, The charge_spin handling is duplicated between _eval_model_spin and _prepare_inputs; extract the shared logic into a single helper (e.g., _extend_spin or prepare_charge_spin) and have both _prepare_inputs and _eval_model_spin call it to produce charge_spin_t (handling numpy/torch, dtype/device, reshape/unsqueeze/expand, default_chg_spin fallback and the ValueError path); update _eval_model_spin to remove its inline block and delegate to the new helper so future dtype or default-fallback changes are made in one place.
1526-1535: 💤 Low valueGate on
dp_am.add_chg_spin_ebdis sound; consider extracting the conditional.The
getattr(dp_am, "add_chg_spin_ebd", False)guard correctly avoids handingcharge_spinto descriptors that don't consume it (everything other than DPA3, per PR objectives). Same conditional is repeated verbatim ineval_fitting_last_layerat lines 1601-1603. A tiny helper (e.g.,_chg_spin_for(dp_am, charge_spin_t)) would DRY this up and make future descriptor additions a one-line change.🤖 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 1526 - 1535, The code repeats the conditional getattr(dp_am, "add_chg_spin_ebd", False) when passing charge_spin to dp_am.descriptor (and later in eval_fitting_last_layer), so extract a small helper (e.g., a function named _chg_spin_for or similar) that takes dp_am and charge_spin_t and returns charge_spin_t if getattr(dp_am, "add_chg_spin_ebd", False) else None; replace the inline conditional in the dp_am.descriptor call and in eval_fitting_last_layer with a call to that helper to DRY the logic and simplify future descriptor additions (reference symbols: dp_am.add_chg_spin_ebd, dp_am.descriptor, eval_fitting_last_layer).
594-635: 💤 Low valueDocument the new
charge_spinparameter ineval's docstring.
charge_spinis now a first-class kwarg on the publicevalAPI, but the docstring still only describescoords,cells,atom_types,atomic,fparam,aparam,**kwargs. Adding a short entry mirroringfparam(shapenframes x 2, optional, falls back todefault_chg_spinwhen the model hasadd_chg_spin_ebd=True) will help library users discover the migration path away from packing charge/spin intofparam.Suggested docstring addition
aparam The atomic parameter. The array should be of size nframes x natoms x dim_aparam. + charge_spin + Per-frame [charge, spin] values for models trained with + ``add_chg_spin_ebd=True``. The array should be of size + ``nframes x 2``. If omitted, the model's ``default_chg_spin`` + (when configured) is used. **kwargs Other parameters🤖 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 594 - 635, Add a docstring entry for the new public parameter charge_spin in the eval method: describe that charge_spin is optional, has shape nframes x 2, and that when omitted the code falls back to default_chg_spin if the model has add_chg_spin_ebd=True (mirror the style of the existing fparam/aparam descriptions); update the parameter list under eval(...) to include this short description so users know to stop packing charge/spin into fparam and where the default comes from (refer to eval, parameter name charge_spin, model flag add_chg_spin_ebd, and default_chg_spin).source/tests/consistent/descriptor/common.py (1)
260-293: 💤 Low value
eval_pd_descriptorsilently dropscharge_spin.The signature accepts
charge_spin, but unlike the other backends (DP/PT/PT_EXPT/JAX/array-api-strict) it is neither converted to a paddle tensor nor forwarded intopd_obj(...). A caller passingcharge_spin=...to this helper will get a result computed without it, with no warning. Per the PR objectives the pd backend is intentionally untouched, so either propagate the kwarg (if pd supports it) or at minimum document this in a docstring/inline comment so future test authors don't assume parity with the other helpers.Suggested clarifying note
def eval_pd_descriptor( self, pd_obj: Any, natoms: np.ndarray, coords: np.ndarray, atype: np.ndarray, box: np.ndarray, mixed_types: bool = False, fparam: np.ndarray | None = None, charge_spin: np.ndarray | None = None, ) -> Any: + # NOTE: the pd backend does not yet consume charge_spin; the kwarg + # is accepted for signature parity with the other eval_*_descriptor + # helpers but is intentionally not forwarded into pd_obj(...). ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pd(🤖 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 `@source/tests/consistent/descriptor/common.py` around lines 260 - 293, The helper eval_pd_descriptor is accepting charge_spin but not forwarding it to pd_obj; either forward it (if pd_obj supports it) by converting charge_spin to a paddle tensor on PD_DEVICE (e.g., paddle.to_tensor(charge_spin).to(PD_DEVICE)) and reshaping to the batch form used for coords/atype (match shape [1, -1, ...] as appropriate) and include charge_spin=... in the pd_obj(...) call, or if pd backend truly must remain unchanged, add a clear inline comment/docstring in eval_pd_descriptor stating that charge_spin is intentionally not handled/forwarded for the pd backend so callers know it will be ignored; update references to eval_pd_descriptor and pd_obj accordingly.
🤖 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.
Inline comments:
In `@deepmd/dpmodel/descriptor/dpa3.py`:
- Around line 419-422: Replace assert-based checks for default_chg_spin and
charge_spin with explicit runtime exceptions: locate the validation blocks in
dpa3.py where self.add_chg_spin_ebd and self.default_chg_spin (and the similar
check for charge_spin around the later block) are validated and change the
assert len(...) == 2 to an explicit conditional that raises a ValueError (or
TypeError) with the same message when the length is not 2; ensure both checks
(the one referencing self.default_chg_spin and the one referencing charge_spin)
are updated so validation still runs under optimized Python.
In `@deepmd/entrypoints/test.py`:
- Around line 642-645: The branch that sets charge_spin uses
test_data.get("find_charge_spin", 1.0) which defaults to enabling the branch and
can cause a KeyError; change the default to 0.0 so missing find_charge_spin
falls through safely (i.e., replace test_data.get("find_charge_spin", 1.0) with
test_data.get("find_charge_spin", 0.0")) and keep the rest of the logic around
dp.has_chg_spin_ebd(), test_data["charge_spin"], and numb_test unchanged.
In `@deepmd/pt/model/atomic_model/base_atomic_model.py`:
- Around line 200-202: The return type of BaseAtomicModel.get_default_chg_spin
is too narrow (torch.Tensor | None) and conflicts with
DPAtomicModel.get_default_chg_spin which returns list[float] | None; update the
type annotation on BaseAtomicModel.get_default_chg_spin to accept both
list[float] and torch.Tensor (e.g., list[float] | torch.Tensor | None) so it
matches DPAtomicModel and other subclasses, and add a note in callers (e.g.,
model_forward where the isinstance guard at lines ~658–661 exists) to convert
list->torch.Tensor via a helper like to_torch_tensor when a torch.Tensor is
required.
In `@deepmd/pt/model/atomic_model/dp_atomic_model.py`:
- Around line 389-413: The new `@torch.jit.export` methods in DPAtomicModel
(has_chg_spin_ebd, get_dim_chg_spin, has_default_chg_spin, get_default_chg_spin)
call into self.descriptor (BaseDescriptor) so add declarations on the
BaseDescriptor class for get_dim_chg_spin() -> int, has_default_chg_spin() ->
bool, and get_default_chg_spin() -> Optional[torch.Tensor] (or torch.Tensor |
None) — either as abstract methods or as concrete defaults (return 0, False,
None) — using TorchScript-compatible types so torch.jit.script() can compile all
code paths; update BaseDescriptor’s import/type hints to include torch.Tensor if
needed.
In `@deepmd/pt/model/descriptor/dpa3.py`:
- Around line 128-131: Replace the runtime assertion in the constructor that
checks default_chg_spin with a proper exception: in the code block handling
default_chg_spin (the conditional around default_chg_spin is not None in
dpa3.py), remove the assert and instead raise a ValueError with the same message
("default_chg_spin must be a list of length 2 [charge, spin].") so validation
cannot be bypassed when Python is run with optimizations; keep the same check
for len(default_chg_spin) == 2 and the existing message to aid debugging.
- Around line 268-277: The get_default_chg_spin method's return annotation uses
PEP 604 union syntax (torch.Tensor | None) which TorchScript/@torch.jit.export
doesn't accept; change the signature of get_default_chg_spin to use
typing.Optional[torch.Tensor] (or typing.Union[torch.Tensor, None]) and add the
necessary import for Optional from typing so the exported method compiles under
TorchScript.
---
Outside diff comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 1117-1141: The positional vs keyword asymmetry can lead to silent
mis-wiring of arguments; update calls to _pt2_runner (in both _eval_model and
_eval_model_spin) to pass charge_spin explicitly as a keyword
(charge_spin=charge_spin_t) instead of as a positional arg (and similarly ensure
ext_spin_t remains positional in the spin path), and add an end-to-end test that
invokes the .pt2 runner with a non-None charge_spin to catch future
argument-reordering regressions.
In `@deepmd/pt_expt/model/spin_ener_model.py`:
- Around line 151-188: SpinModel.forward_common_lower_exportable currently omits
the charge_spin parameter causing a signature mismatch when
SpinEnerModel.forward_lower_exportable passes charge_spin through; update
SpinModel.forward_common_lower_exportable to add charge_spin to its signature,
include charge_spin as the 7th positional parameter in the inner traced wrapper
(the inner fn that currently accepts extended_coord, extended_atype,
extended_spin, nlist, mapping, fparam, aparam, ...) and ensure charge_spin is
forwarded into the traced(...) call and through any make_fx_kwargs plumbing (so
it is not swallowed by **make_fx_kwargs). Reference
SpinModel.forward_common_lower_exportable,
SpinEnerModel.forward_lower_exportable, the traced object and inner fn to find
and make the change.
In `@deepmd/pt/model/descriptor/dpa3.py`:
- Around line 206-211: The charge embedding chg_embedding currently constructs
TypeEmbedNet(200, ...) while inputs are shifted by +100, causing an off-by-one
(charge=100 -> index 200) and IndexError; change the embedding size to 201 so
indices 0..200 cover charges -100..100 and ensure the forward mapping that adds
100 uses that new range. Also audit the spin embedding (TypeEmbedNet(100, ...)
and any uses that index it) and add a clamp or assert in the forward method to
enforce valid spin and charge ranges (or expand the embedding size if the
intended inclusive range is larger); apply the same fix to the other occurrence
around the spin/charge embedding block referenced later (lines ~586-597).
---
Nitpick comments:
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 301-334: The call path forwards charge_spin without type
normalization, so modify call_common and call_common_lower to include
charge_spin in the inputs processed by _input_type_cast (or explicitly cast
charge_spin to the same float dtype as returned by _input_type_cast) so that
charge_spin is converted to input_prec before calling model_call_from_call_lower
or passing into the lower-level descriptor; update the call sites: in
call_common (before invoking model_call_from_call_lower) and in
call_common_lower (at the start of the function) to use the normalized/converted
charge_spin variable returned or produced by _input_type_cast, ensuring all
arrays (coord/box/fparam/aparam/charge_spin) share the same dtype.
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 1273-1297: The charge_spin handling is duplicated between
_eval_model_spin and _prepare_inputs; extract the shared logic into a single
helper (e.g., _extend_spin or prepare_charge_spin) and have both _prepare_inputs
and _eval_model_spin call it to produce charge_spin_t (handling numpy/torch,
dtype/device, reshape/unsqueeze/expand, default_chg_spin fallback and the
ValueError path); update _eval_model_spin to remove its inline block and
delegate to the new helper so future dtype or default-fallback changes are made
in one place.
- Around line 1526-1535: The code repeats the conditional getattr(dp_am,
"add_chg_spin_ebd", False) when passing charge_spin to dp_am.descriptor (and
later in eval_fitting_last_layer), so extract a small helper (e.g., a function
named _chg_spin_for or similar) that takes dp_am and charge_spin_t and returns
charge_spin_t if getattr(dp_am, "add_chg_spin_ebd", False) else None; replace
the inline conditional in the dp_am.descriptor call and in
eval_fitting_last_layer with a call to that helper to DRY the logic and simplify
future descriptor additions (reference symbols: dp_am.add_chg_spin_ebd,
dp_am.descriptor, eval_fitting_last_layer).
- Around line 594-635: Add a docstring entry for the new public parameter
charge_spin in the eval method: describe that charge_spin is optional, has shape
nframes x 2, and that when omitted the code falls back to default_chg_spin if
the model has add_chg_spin_ebd=True (mirror the style of the existing
fparam/aparam descriptions); update the parameter list under eval(...) to
include this short description so users know to stop packing charge/spin into
fparam and where the default comes from (refer to eval, parameter name
charge_spin, model flag add_chg_spin_ebd, and default_chg_spin).
In `@source/tests/consistent/descriptor/common.py`:
- Around line 260-293: The helper eval_pd_descriptor is accepting charge_spin
but not forwarding it to pd_obj; either forward it (if pd_obj supports it) by
converting charge_spin to a paddle tensor on PD_DEVICE (e.g.,
paddle.to_tensor(charge_spin).to(PD_DEVICE)) and reshaping to the batch form
used for coords/atype (match shape [1, -1, ...] as appropriate) and include
charge_spin=... in the pd_obj(...) call, or if pd backend truly must remain
unchanged, add a clear inline comment/docstring in eval_pd_descriptor stating
that charge_spin is intentionally not handled/forwarded for the pd backend so
callers know it will be ignored; update references to eval_pd_descriptor and
pd_obj accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 30459192-fb59-447b-9f10-92019ad210db
📒 Files selected for processing (60)
deepmd/calculator.pydeepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/atomic_model/linear_atomic_model.pydeepmd/dpmodel/atomic_model/make_base_atomic_model.pydeepmd/dpmodel/atomic_model/pairtab_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/dpa2.pydeepmd/dpmodel/descriptor/dpa3.pydeepmd/dpmodel/descriptor/hybrid.pydeepmd/dpmodel/descriptor/make_base_descriptor.pydeepmd/dpmodel/descriptor/se_e2_a.pydeepmd/dpmodel/descriptor/se_r.pydeepmd/dpmodel/descriptor/se_t.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/model/ener_model.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/utils/batch.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/dpmodel/utils/stat.pydeepmd/entrypoints/test.pydeepmd/infer/deep_eval.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/atomic_model/pairtab_atomic_model.pydeepmd/pt/model/descriptor/dpa1.pydeepmd/pt/model/descriptor/dpa2.pydeepmd/pt/model/descriptor/dpa3.pydeepmd/pt/model/descriptor/hybrid.pydeepmd/pt/model/descriptor/se_a.pydeepmd/pt/model/descriptor/se_r.pydeepmd/pt/model/descriptor/se_t.pydeepmd/pt/model/descriptor/se_t_tebd.pydeepmd/pt/model/model/ener_model.pydeepmd/pt/model/model/make_model.pydeepmd/pt/train/training.pydeepmd/pt/train/wrapper.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/dipole_model.pydeepmd/pt_expt/model/dos_model.pydeepmd/pt_expt/model/dp_linear_model.pydeepmd/pt_expt/model/dp_zbl_model.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/polar_model.pydeepmd/pt_expt/model/property_model.pydeepmd/pt_expt/model/spin_ener_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/wrapper.pydeepmd/utils/argcheck.pysource/tests/consistent/descriptor/common.pysource/tests/consistent/descriptor/test_dpa3.pysource/tests/consistent/model/test_ener.pysource/tests/pt/model/test_dpa3.pysource/tests/pt_expt/descriptor/test_dpa3.pysource/tests/universal/dpmodel/descriptor/test_descriptor.pysource/tests/universal/dpmodel/model/test_model.py
Comment on lines
+419
to
+422
| if self.add_chg_spin_ebd and self.default_chg_spin is not None: | ||
| assert len(self.default_chg_spin) == 2, ( | ||
| "default_chg_spin must have exactly 2 values [charge, spin]" | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="$(fd -p 'dpa3.py' | head -n 1)"
echo "Inspecting: $FILE"
nl -ba "$FILE" | sed -n '400,710p' | rg -n 'assert|default_chg_spin|charge_spin'Repository: deepmodeling/deepmd-kit
Length of output: 172
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="$(fd -p 'dpa3.py' | head -n 1)"
echo "Inspecting: $FILE"
echo "=== Lines 415-425 (around first assert) ==="
cat -n "$FILE" | sed -n '415,425p'
echo ""
echo "=== Lines 685-695 (around second assert) ==="
cat -n "$FILE" | sed -n '685,695p'Repository: deepmodeling/deepmd-kit
Length of output: 1327
Replace assert-based input checks with explicit exceptions.
Lines 420-422 and 690-692 rely on assert for runtime validation. In optimized Python runs (-O), those checks are removed, so invalid default_chg_spin/charge_spin can bypass validation and fail later in less controlled ways.
Proposed fix
self.add_chg_spin_ebd = add_chg_spin_ebd
self.default_chg_spin = default_chg_spin
if self.add_chg_spin_ebd and self.default_chg_spin is not None:
- assert len(self.default_chg_spin) == 2, (
- "default_chg_spin must have exactly 2 values [charge, spin]"
- )
+ if len(self.default_chg_spin) != 2:
+ raise ValueError(
+ "default_chg_spin must have exactly 2 values [charge, spin]"
+ )
...
if self.add_chg_spin_ebd:
- assert charge_spin is not None
- assert self.chg_embedding is not None
- assert self.spin_embedding is not None
+ if charge_spin is None:
+ raise ValueError(
+ "charge_spin must be provided when add_chg_spin_ebd is enabled."
+ )
+ if self.chg_embedding is None or self.spin_embedding is None:
+ raise RuntimeError("charge/spin embeddings are not initialized.")🤖 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/dpmodel/descriptor/dpa3.py` around lines 419 - 422, Replace
assert-based checks for default_chg_spin and charge_spin with explicit runtime
exceptions: locate the validation blocks in dpa3.py where self.add_chg_spin_ebd
and self.default_chg_spin (and the similar check for charge_spin around the
later block) are validated and change the assert len(...) == 2 to an explicit
conditional that raises a ValueError (or TypeError) with the same message when
the length is not 2; ensure both checks (the one referencing
self.default_chg_spin and the one referencing charge_spin) are updated so
validation still runs under optimized Python.
Comment on lines
+642
to
+645
| if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 1.0) != 0.0: | ||
| charge_spin = test_data["charge_spin"][:numb_test] | ||
| else: | ||
| charge_spin = None |
Contributor
There was a problem hiding this comment.
Unsafe default for find_charge_spin — should be 0.0, not 1.0.
test_data.get("find_charge_spin", 1.0) differs from the established fparam pattern (line 634 uses direct access test_data["find_fparam"]). If find_charge_spin is absent from test_data (e.g., in an unexpected LMDB layout), a default of 1.0 makes the condition True, which then unconditionally accesses test_data["charge_spin"] and raises KeyError. A default of 0.0 falls through safely to charge_spin = None.
🛡️ Proposed fix
- if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 1.0) != 0.0:
+ if dp.has_chg_spin_ebd() and test_data.get("find_charge_spin", 0.0) != 0.0:🤖 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/entrypoints/test.py` around lines 642 - 645, The branch that sets
charge_spin uses test_data.get("find_charge_spin", 1.0) which defaults to
enabling the branch and can cause a KeyError; change the default to 0.0 so
missing find_charge_spin falls through safely (i.e., replace
test_data.get("find_charge_spin", 1.0) with test_data.get("find_charge_spin",
0.0")) and keep the rest of the logic around dp.has_chg_spin_ebd(),
test_data["charge_spin"], and numb_test unchanged.
Comment on lines
+200
to
+202
| def get_default_chg_spin(self) -> torch.Tensor | None: | ||
| """Get the default charge_spin values.""" | ||
| return None |
Contributor
There was a problem hiding this comment.
Return type annotation conflicts with dpmodel subclass implementations
get_default_chg_spin() is declared as returning torch.Tensor | None in PT BaseAtomicModel, but its concrete dpmodel counterpart (DPAtomicModel.get_default_chg_spin() in deepmd/dpmodel/atomic_model/dp_atomic_model.py) returns list[float] | None. PT's DPAtomicModel inherits from the dpmodel version; its actual return type does not conform to the torch.Tensor | None contract declared here. The isinstance guard added at lines 658–661 acts as a runtime workaround for the model_forward path but does not fix the annotation contract for all callers.
🛠️ Suggested fix — align return type with dpmodel
- def get_default_chg_spin(self) -> torch.Tensor | None:
+ def get_default_chg_spin(self) -> list[float] | None:
"""Get the default charge_spin values."""
return NoneAny call sites that need a torch.Tensor should convert explicitly (e.g., via to_torch_tensor).
🤖 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/model/atomic_model/base_atomic_model.py` around lines 200 - 202,
The return type of BaseAtomicModel.get_default_chg_spin is too narrow
(torch.Tensor | None) and conflicts with DPAtomicModel.get_default_chg_spin
which returns list[float] | None; update the type annotation on
BaseAtomicModel.get_default_chg_spin to accept both list[float] and torch.Tensor
(e.g., list[float] | torch.Tensor | None) so it matches DPAtomicModel and other
subclasses, and add a note in callers (e.g., model_forward where the isinstance
guard at lines ~658–661 exists) to convert list->torch.Tensor via a helper like
to_torch_tensor when a torch.Tensor is required.
Comment on lines
+389
to
+413
| @torch.jit.export | ||
| def has_chg_spin_ebd(self) -> bool: | ||
| """Check if the model has charge spin embedding.""" | ||
| return self.add_chg_spin_ebd | ||
|
|
||
| @torch.jit.export | ||
| def get_dim_chg_spin(self) -> int: | ||
| """Get the dimension of charge_spin input.""" | ||
| if self.add_chg_spin_ebd: | ||
| return self.descriptor.get_dim_chg_spin() | ||
| return 0 | ||
|
|
||
| @torch.jit.export | ||
| def has_default_chg_spin(self) -> bool: | ||
| """Check if the model has default charge_spin values.""" | ||
| if self.add_chg_spin_ebd: | ||
| return self.descriptor.has_default_chg_spin() | ||
| return False | ||
|
|
||
| @torch.jit.export | ||
| def get_default_chg_spin(self) -> torch.Tensor | None: | ||
| """Get the default charge_spin values as a tensor.""" | ||
| if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): | ||
| return self.descriptor.get_default_chg_spin() | ||
| return None |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify that BaseDescriptor in the PT backend declares the charge_spin query methods
# that DPAtomicModel's `@torch.jit.export` methods call.
# Expect: all four method names to appear in base_descriptor.py.
rg -n "def get_dim_chg_spin\|def has_default_chg_spin\|def get_default_chg_spin\|def has_chg_spin_ebd" \
deepmd/pt/model/descriptor/base_descriptor.py
# Also scan dpmodel BaseDescriptor for completeness (used by dpmodel backend)
rg -n "def get_dim_chg_spin\|def has_default_chg_spin\|def get_default_chg_spin\|def has_chg_spin_ebd" \
deepmd/dpmodel/descriptor/base_descriptor.py 2>/dev/null || trueRepository: deepmodeling/deepmd-kit
Length of output: 49
BaseDescriptor must declare the charge-spin query methods called by new @torch.jit.export methods.
All four new @torch.jit.export methods — has_chg_spin_ebd, get_dim_chg_spin, has_default_chg_spin, and get_default_chg_spin — call methods on self.descriptor (typed as BaseDescriptor). TorchScript compiles all code paths statically, so the if self.add_chg_spin_ebd: guards do not prevent compilation. BaseDescriptor must declare get_dim_chg_spin(), has_default_chg_spin(), and get_default_chg_spin() (at minimum as abstract methods or with default implementations). Without these declarations, torch.jit.script() will fail for any model that wraps a DPAtomicModel, breaking the inference path.
🤖 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/model/atomic_model/dp_atomic_model.py` around lines 389 - 413, The
new `@torch.jit.export` methods in DPAtomicModel (has_chg_spin_ebd,
get_dim_chg_spin, has_default_chg_spin, get_default_chg_spin) call into
self.descriptor (BaseDescriptor) so add declarations on the BaseDescriptor class
for get_dim_chg_spin() -> int, has_default_chg_spin() -> bool, and
get_default_chg_spin() -> Optional[torch.Tensor] (or torch.Tensor | None) —
either as abstract methods or as concrete defaults (return 0, False, None) —
using TorchScript-compatible types so torch.jit.script() can compile all code
paths; update BaseDescriptor’s import/type hints to include torch.Tensor if
needed.
Comment on lines
+128
to
+131
| if default_chg_spin is not None: | ||
| assert len(default_chg_spin) == 2, ( | ||
| "default_chg_spin must be a list of length 2 [charge, spin]." | ||
| ) |
Contributor
There was a problem hiding this comment.
Prefer raise ValueError over assert for public constructor validation.
assert can be silently disabled when Python is run with -O, removing the only guard against a malformed default_chg_spin list.
🛡️ Proposed fix
- if default_chg_spin is not None:
- assert len(default_chg_spin) == 2, (
- "default_chg_spin must be a list of length 2 [charge, spin]."
- )
+ if default_chg_spin is not None and len(default_chg_spin) != 2:
+ raise ValueError(
+ "default_chg_spin must be a list of length 2 [charge, spin]."
+ )🤖 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/model/descriptor/dpa3.py` around lines 128 - 131, Replace the
runtime assertion in the constructor that checks default_chg_spin with a proper
exception: in the code block handling default_chg_spin (the conditional around
default_chg_spin is not None in dpa3.py), remove the assert and instead raise a
ValueError with the same message ("default_chg_spin must be a list of length 2
[charge, spin].") so validation cannot be bypassed when Python is run with
optimizations; keep the same check for len(default_chg_spin) == 2 and the
existing message to aid debugging.
Comment on lines
+268
to
+277
| @torch.jit.export | ||
| def get_default_chg_spin(self) -> torch.Tensor | None: | ||
| """Get the default charge_spin values as a tensor.""" | ||
| if self.default_chg_spin is None: | ||
| return None | ||
| return torch.tensor( | ||
| self.default_chg_spin, | ||
| dtype=self.prec, | ||
| device=env.DEVICE, | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
Does PyTorch TorchScript support PEP 604 union syntax (e.g., torch.Tensor | None) in @torch.jit.export method return type annotations in PyTorch 2.11?
💡 Result:
No, PyTorch TorchScript does not support PEP 604 union syntax (e.g., torch.Tensor | None) in @torch.jit.export method return type annotations in PyTorch 2.11. TorchScript is deprecated and in maintenance mode [1][2], with limited support for new Python features. While basic Union types from typing have been supported since around PyTorch 1.12 [3], and PEP 604 | syntax was requested [4][5] and partially added via PR #109293 around PyTorch 2.2 [4], issues persist in 2025 with union usage, such as unpacking unions [6] and optimization flags [7]. Specific errors like "Expression of type | cannot be used in a type expression" were reported for torch.jit.script [5], and TorchScript's type parsing does not fully handle | in annotations reliably [4]. Examples using torch.Tensor | None or similar in @torch.jit.export methods fail compilation, as evidenced by ongoing GitHub issues post-2.2 [6][7]. Use Optional[torch.Tensor] or Union[torch.Tensor, None] instead for compatibility. PyTorch recommends migrating to torch.export or torch.compile [2].
Citations:
- 1: https://github.com/pytorch/pytorch/releases/tag/v2.11.0
- 2: https://pytorch.org/blog/pytorch-2-11-release-blog/
- 3: Support Union in TorchScript pytorch/pytorch#64234
- 4: Torchscript should support union type by "|" introduced in python3.10. pytorch/pytorch#114755
- 5: Python 3.10 Union operator
|is not supported in jit script pytorch/pytorch#101777 - 6: TorchScript: Union cannot be used as a tuple pytorch/pytorch#123168
- 7: Union type raise error when running python with argument "-O" for torch script. pytorch/pytorch#148447
🏁 Script executed:
head -280 deepmd/pt/model/descriptor/dpa3.py | tail -20Repository: deepmodeling/deepmd-kit
Length of output: 780
Change return type annotation from torch.Tensor | None to Optional[torch.Tensor].
The @torch.jit.export decorator requires Optional[torch.Tensor] from typing for the return type; PyTorch TorchScript does not support PEP 604 union syntax (X | None) in exported method signatures as of PyTorch 2.11. Replace with Optional[torch.Tensor] or Union[torch.Tensor, None] to ensure TorchScript compilation succeeds.
🤖 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/model/descriptor/dpa3.py` around lines 268 - 277, The
get_default_chg_spin method's return annotation uses PEP 604 union syntax
(torch.Tensor | None) which TorchScript/@torch.jit.export doesn't accept; change
the signature of get_default_chg_spin to use typing.Optional[torch.Tensor] (or
typing.Union[torch.Tensor, None]) and add the necessary import for Optional from
typing so the exported method compiles under TorchScript.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
add_chg_spin_ebd=Truepreviously hijackedfparamto smuggle the[charge, spin] scalars into DPA3, forcing users to set
numb_fparam=2on the fitting net and blocking real frame parameters from coexisting
with charge/spin. This PR plumbs
charge_spin: Tensor | Noneas afirst-class kwarg through every forward chain and adds an optional
default_chg_spinfallback on the DPA3 descriptor.Backends covered: pt, dpmodel, pt_expt. The pd backend is left untouched.
The C/C++/LAMMPS layer is unchanged.
Forward chain
Calculator / deep_eval / dp test / lmdb_data / training.get_data->
wrapper.forward->
ener_model.forward / forward_lower->
make_model.forward_common / forward_common_lower->
base_atomic_model.forward_common_atomic->
dp_atomic_model.forward_atomic# default_chg_spin fallback here->
descriptor.forward# only DPA3 consumes itAll other descriptors (se_e2_a, se_r, se_t, se_t_tebd, dpa1, dpa2,
hybrid) only forward the kwarg through their signatures.
New API surface
On
BaseAtomicModeland the wrapped model:has_chg_spin_ebd() -> boolget_dim_chg_spin() -> int# 2 for DPA3, else 0has_default_chg_spin() -> boolget_default_chg_spin() -> list[float] | Tensor | NoneDPA3 descriptor gains a
default_chg_spin: list[float] | None = Noneconstructor arg (length 2, validated; round-trips through
serialize).descrpt_dpa3_argsexposes the matchingArgumentand theadd_chg_spin_ebddoc no longer references fparam.Training data
charge_spinis registered as aDataRequirementItem(ndof=2, atomic=False, must=not has_default_cs, default=cs_default). Theget_datapath drops it (along with fparam) on frames wherefind_charge_spin == 0, so missing per-frame data falls back todefault_chg_spinwhen one is configured.pt_expt specifics
forward_common_atomic,forward_common_lower_exportable, themake_fx-traced innerfn,_trace_and_compile, and all wrappingenergy/spin/dipole/dos/polar/property/dp_linear/dp_zbl model variants
gained a
charge_spinarg in lockstep so the export and inductor-compiled paths keep matching signatures.
deep_evalno longer reusesfparamfor charge/spin — it constructscharge_spin_t(with themetadata default-fallback) and passes it explicitly.
Tests
Three
cs_modecases are exercised everywhere it matters:no_chg_spin,explicit_chg_spin,default_chg_spin.source/tests/pt/model/test_dpa3.py::test_consistency)rewritten over the three modes; default mode also asserts that the
default-fallback descriptor matches an explicit
[5,1]peer.source/tests/pt_expt/descriptor/test_dpa3.py) gainstest_consistency_chg_spincovering explicit and default modesagainst dpmodel.
DescriptorParamDPA3learnsdefault_chg_spin,parametrize gains
(None, [5.0, 1.0]), and theadd_chg_spin_ebdskip rule intest_model.pyis replaced —the universal driver does not feed
charge_spin, so chg_spin runsrely on the
default_chg_spinfallback. 622 DPA3 model cases pass.descriptor/common.pythreadscharge_spinthrough every
eval_*(pd ignores it).test_dpa3.pyswapsself.fparamforself.charge_spin.test_ener.py:: TestEnerChgSpinEbdFparamis reparametrized over the three modesand no longer touches
numb_fparam/default_fparam.Smoke
examples/water/dpa3 dp --pt train input_torch_dynamic.json --skip-neighbor-statruns to batch 600 with monotonically decreasingloss.
Test plan
Summary by CodeRabbit
Release Notes
New Features
Tests