feat(dpa3): decouple charge_spin from fparam#5431
Open
iProzd wants to merge 9 commits intodeepmodeling:masterfrom
Open
feat(dpa3): decouple charge_spin from fparam#5431iProzd wants to merge 9 commits intodeepmodeling:masterfrom
iProzd wants to merge 9 commits intodeepmodeling:masterfrom
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
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
|
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:
📝 WalkthroughWalkthroughThis PR adds dedicated ChangesCharge-Spin Input Threading
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~35 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![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
+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.
wanghan-iapcm
requested changes
May 6, 2026
| use_loc_mapping: bool = True, | ||
| type_map: list[str] | None = None, | ||
| add_chg_spin_ebd: bool = False, | ||
| default_chg_spin: list[float] | None = None, |
Collaborator
There was a problem hiding this comment.
default_chg_spin could be given [0.0, 0.0].
wanghan-iapcm
requested changes
May 6, 2026
Collaborator
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
- this is a breaking change. when trying to use old model with
default_fparaminjecting to the fitting net, an error should be raised. - dipole_model.py, dos_model.py, dp_zbl_model.py, polar_model.py, property_model.py all show large diffs (+155/−142 etc.) where the substantive change is ~13 lines of charge_spin plumbing. The remainder is line-ending / trailing-whitespace re-formatting that makes review and git blame hard.
- deepmd/pt/train/training.py:1909 and deepmd/pt_expt/train/training.py:141:
has_chg_spin_ebdis now defined on the base atomic model — so the False default can never be returned (it's always a method). Thecallable()check is dead-code-defensive. Simplify to a direct call.
has_chg_spin_ebd = getattr(_model, "has_chg_spin_ebd", False)
if callable(has_chg_spin_ebd):
has_chg_spin_ebd = has_chg_spin_ebd()
wanghan-iapcm
requested changes
May 6, 2026
| atol=atol, | ||
| ) | ||
|
|
||
| # default_chg_spin should match explicit when value is the same. |
Collaborator
There was a problem hiding this comment.
I would recommend moving the consistency check to dpmodel backend, as pt will be deprecated in the future.
njzjz-bot
suggested changes
May 6, 2026
Contributor
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
I agree with the direction of making charge_spin a first-class input instead of overloading fparam, but I do not think this PR is mergeable yet.
Blocking points from my side:
- CI is currently heavily failing/cancelled across Python and C++ jobs. Please make the test suite green before this can be reviewed for merge.
- This changes the public inference/training interface and the data schema (
charge_spinvs the previousfparamconvention for DPA3 charge/spin). Please document the migration path clearly: what happens to existing DPA3 charge/spin models or datasets that only providefparam, whether they are intentionally unsupported, and what users should change. - The required-data/default path needs more coverage. In particular, please add/confirm tests for all three modes: no charge-spin embedding, explicit
charge_spin, anddefault_chg_spin, across training and inference (deepmd test/ASE/eval path if applicable). - The new API plumbing is broad (dpmodel, pt, pt_expt, exported/compiled paths). Please make sure exported PT2 / AOTInductor call signatures remain backward compatible for models without charge-spin, and that missing
charge_spinraises a clear error only whenadd_chg_spin_ebd=Trueand no default is present. charge_spinis being reshaped to(nframes, 2)in several infer paths. Please validate shape early with a clear error instead of relying on reshape/assert behavior, so accidental per-atom or flattened inputs fail understandably.
Once CI is green and the compatibility/migration story is explicit, I am happy to take another pass.
Reviewed by OpenClaw 2026.4.22 (00bd2cf) (model: gpt-5.5).
Authored by OpenClaw (model: gpt-5.5)
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/pt_expt/model/dipole_model.py`:
- Around line 47-48: You added a new parameter charge_spin before
do_atomic_virial which breaks positional-call compatibility; fix all signatures
that introduced charge_spin (the constructor and the other two methods that
currently list charge_spin then do_atomic_virial) by either placing
do_atomic_virial before charge_spin or by making charge_spin a keyword-only
parameter (e.g., keep do_atomic_virial as the last positional arg and move
charge_spin after a * so callers passing do_atomic_virial positionally continue
to work); update every function/method signature that references charge_spin and
do_atomic_virial 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: d1f1905d-b978-498c-ab0d-80b576ace353
📒 Files selected for processing (5)
deepmd/pt_expt/model/dipole_model.pydeepmd/pt_expt/model/dos_model.pydeepmd/pt_expt/model/dp_zbl_model.pydeepmd/pt_expt/model/polar_model.pydeepmd/pt_expt/model/property_model.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/pt_expt/model/polar_model.py
- deepmd/pt_expt/model/property_model.py
Comment on lines
+47
to
48
| charge_spin: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, |
Contributor
There was a problem hiding this comment.
Preserve positional-call compatibility when adding charge_spin
At Line 47, Line 80, and Line 132, inserting charge_spin before do_atomic_virial changes positional binding and can silently break existing callers passing do_atomic_virial positionally.
Suggested compatibility-safe adjustment
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,
+ charge_spin: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]: 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,
- charge_spin: torch.Tensor | None = None,
do_atomic_virial: bool = False,
+ charge_spin: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]: def forward_lower_exportable(
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,
- charge_spin: torch.Tensor | None = None,
do_atomic_virial: bool = False,
+ charge_spin: torch.Tensor | None = None,
**make_fx_kwargs: Any,
) -> torch.nn.Module:Also applies to: 80-81, 132-133
🤖 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/dipole_model.py` around lines 47 - 48, You added a new
parameter charge_spin before do_atomic_virial which breaks positional-call
compatibility; fix all signatures that introduced charge_spin (the constructor
and the other two methods that currently list charge_spin then do_atomic_virial)
by either placing do_atomic_virial before charge_spin or by making charge_spin a
keyword-only parameter (e.g., keep do_atomic_virial as the last positional arg
and move charge_spin after a * so callers passing do_atomic_virial positionally
continue to work); update every function/method signature that references
charge_spin and do_atomic_virial accordingly.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
deepmd/jax/model/dp_zbl_model.py (1)
31-54:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftSame downstream
TypeErrorasdp_model.py— root cause is the missingcharge_spinkwarg on the JAX atomic model side.
DPZBLLinearEnergyAtomicModel.forward_common_atomicis also a JAX atomic model withoutcharge_spinin its signature. Once the JAX atomic models are updated (seebase_model.pycomment), these changes will be unblocked.🤖 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/dp_zbl_model.py` around lines 31 - 54, The DPZBL JAX atomic wrapper forwards to forward_common_atomic but the underlying JAX atomic model DPZBLLinearEnergyAtomicModel is missing the charge_spin parameter; update DPZBLLinearEnergyAtomicModel.forward_common_atomic (and any JAX atomic model signatures in this module referenced from base_model.py) to accept charge_spin: jnp.ndarray | None = None and pass it through to the internal forward_common_atomic call so the wrapper's charge_spin kwarg is honored.deepmd/jax/model/base_model.py (2)
62-161:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
charge_spincaptured as closure variable ineval_output/eval_ce— not batched byjax.vmap, causing shape mismatch fornf > 1.
fparamandaparamare explicit positional arguments ofeval_output(andeval_ce), sojax.vmapcorrectly reduces them from[nf, d]to[d]per frame, after whichfparam[None, ...]restores the batch dimension.charge_spinis instead captured from the outer scope at lines 80 and 143, so each vmapped call receives the full[nf, 2]tensor rather than a per-frame[2]. When passed ascharge_spin=charge_spin(without[None, ...]), the atomic model sees shape[nf, 2]instead of the expected[1, 2], leading to a shape mismatch at runtime once the atomic-model signature is fixed.The same issue applies to the
jax.hessiancall (lines 103–110), which also invokeseval_outputwithout passingcharge_spin.💡 Proposed fix
def eval_output( cc_ext: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, mapping: jnp.ndarray | None, fparam: jnp.ndarray | None, aparam: jnp.ndarray | None, + charge_spin_: jnp.ndarray | None, *, _kk: str = kk, _atom_axis: int = atom_axis, ) -> jnp.ndarray: atomic_ret = self.atomic_model.forward_common_atomic( cc_ext[None, ...], extended_atype[None, ...], nlist[None, ...], mapping=mapping[None, ...] if mapping is not None else None, fparam=fparam[None, ...] if fparam is not None else None, aparam=aparam[None, ...] if aparam is not None else None, - charge_spin=charge_spin, + charge_spin=charge_spin_[None, ...] if charge_spin_ is not None else None, ) return jnp.sum(atomic_ret[_kk][0], axis=_atom_axis) ff = -jax.vmap(jax.jacrev(eval_output, argnums=0))( extended_coord, extended_atype, nlist, mapping, fparam, aparam, + charge_spin, ) ... if vdef.r_hessian: hessian = jax.vmap(jax.hessian(eval_output, argnums=0))( extended_coord, extended_atype, nlist, mapping, fparam, aparam, + charge_spin, )Apply the same treatment to
eval_ceand itsjax.vmap(jax.jacrev(eval_ce))call (lines 124–161).🤖 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/base_model.py` around lines 62 - 161, The closure captures charge_spin causing vmapped calls to receive the full [nf,2] tensor; make charge_spin an explicit argument to eval_output and eval_ce (e.g., add charge_spin: jnp.ndarray to their signatures) and update all jax.vmap/jax.jacrev/jax.hessian invocations that call these functions to pass the per-frame charge_spin (so vmapping will slice [nf,2] -> [2] per invocation); inside each function keep using charge_spin as before (atomic_model.forward_common_atomic(..., charge_spin=charge_spin)) so the atomic model receives a [1,2] batch when you wrap with [None,...] as currently done.
31-38:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftJAX atomic models must accept
charge_spinparameter to unblock pipeline failures.The call to
self.atomic_model.forward_common_atomic(..., charge_spin=charge_spin)at line 38 will raiseTypeErrorbecause the JAX atomic model implementations (DPAtomicModel,DPZBLLinearEnergyAtomicModel, andPairTabAtomicModel) do not declare thecharge_spinparameter in theirforward_common_atomicsignatures, while the dpmodel base implementation does. Addcharge_spin: jnp.ndarray | None = Noneto the signatures offorward_common_atomicin all affected JAX atomic model classes and forward the parameter to their parent implementations.🤖 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/base_model.py` around lines 31 - 38, The JAX atomic model classes (DPAtomicModel, DPZBLLinearEnergyAtomicModel, PairTabAtomicModel) lack the charge_spin parameter in their forward_common_atomic signatures but are called with charge_spin from base_model; update each class's forward_common_atomic method to add the parameter signature charge_spin: jnp.ndarray | None = None and pass it along when calling the parent/base implementation (i.e., include charge_spin in super().forward_common_atomic(...) or in the internal call that delegates to the dpmodel base), ensuring the parameter is accepted and forwarded to prevent the TypeError.deepmd/jax/model/dp_model.py (1)
49-72:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftAll JAX atomic models lack
charge_spinparameter—update theirforward_common_atomicsignatures.The error at line 61 occurs because
base_model.forward_common_atomic()passescharge_spin=charge_spintoself.atomic_model.forward_common_atomic(), but the JAX atomic model implementations (pairtab_atomic_model, linear_atomic_model, dp_atomic_model) do not accept this parameter. Update all three atomic model signatures to includecharge_spin: jnp.ndarray | None = None(and forward it to their parentsuper().forward_common_atomic()calls if needed).🤖 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/dp_model.py` around lines 49 - 72, The JAX atomic model methods (pairtab_atomic_model.forward_common_atomic, linear_atomic_model.forward_common_atomic, dp_atomic_model.forward_common_atomic) must accept the missing parameter charge_spin: jnp.ndarray | None = None and forward it when calling their parent/super implementation; update each method signature to add the charge_spin parameter and propagate it in the call to super().forward_common_atomic (or return wrapper) so base_model.forward_common_atomic's charge_spin argument is accepted end-to-end.
🤖 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/pt/model/descriptor/hybrid.py`:
- Around line 348-350: The call to descrpt in the hybrid descriptor path only
forwards charge_spin and drops other supported keyword arguments (e.g.,
comm_dict, fparam), silently changing sub-descriptor behavior; update the call
site around the descrpt invocation so it forwards the remaining kwargs—either
pass comm_dict=comm_dict and fparam=fparam (ensuring those names exist in scope)
or accept and forward **kwargs from the hybrid descriptor function into
descrpt—so all supported options reach the sub-descriptors.
- Around line 102-112: The three hardcoded accessors (get_dim_chg_spin,
has_default_chg_spin, get_default_chg_spin) must aggregate capabilities from the
child descriptors instead of returning 0/False/None; update get_dim_chg_spin to
iterate your child list (e.g., self.descriptors or whatever container the hybrid
holds) and return the sum of each child's get_dim_chg_spin(), update
has_default_chg_spin to return True if any child.has_default_chg_spin() is True,
and update get_default_chg_spin to return a combined default (e.g., concatenate
child.get_default_chg_spin() values in descriptor order) or None if no child
provides a default—handle None values when concatenating and keep ordering
consistent with how inputs are constructed.
---
Outside diff comments:
In `@deepmd/jax/model/base_model.py`:
- Around line 62-161: The closure captures charge_spin causing vmapped calls to
receive the full [nf,2] tensor; make charge_spin an explicit argument to
eval_output and eval_ce (e.g., add charge_spin: jnp.ndarray to their signatures)
and update all jax.vmap/jax.jacrev/jax.hessian invocations that call these
functions to pass the per-frame charge_spin (so vmapping will slice [nf,2] ->
[2] per invocation); inside each function keep using charge_spin as before
(atomic_model.forward_common_atomic(..., charge_spin=charge_spin)) so the atomic
model receives a [1,2] batch when you wrap with [None,...] as currently done.
- Around line 31-38: The JAX atomic model classes (DPAtomicModel,
DPZBLLinearEnergyAtomicModel, PairTabAtomicModel) lack the charge_spin parameter
in their forward_common_atomic signatures but are called with charge_spin from
base_model; update each class's forward_common_atomic method to add the
parameter signature charge_spin: jnp.ndarray | None = None and pass it along
when calling the parent/base implementation (i.e., include charge_spin in
super().forward_common_atomic(...) or in the internal call that delegates to the
dpmodel base), ensuring the parameter is accepted and forwarded to prevent the
TypeError.
In `@deepmd/jax/model/dp_model.py`:
- Around line 49-72: The JAX atomic model methods
(pairtab_atomic_model.forward_common_atomic,
linear_atomic_model.forward_common_atomic,
dp_atomic_model.forward_common_atomic) must accept the missing parameter
charge_spin: jnp.ndarray | None = None and forward it when calling their
parent/super implementation; update each method signature to add the charge_spin
parameter and propagate it in the call to super().forward_common_atomic (or
return wrapper) so base_model.forward_common_atomic's charge_spin argument is
accepted end-to-end.
In `@deepmd/jax/model/dp_zbl_model.py`:
- Around line 31-54: The DPZBL JAX atomic wrapper forwards to
forward_common_atomic but the underlying JAX atomic model
DPZBLLinearEnergyAtomicModel is missing the charge_spin parameter; update
DPZBLLinearEnergyAtomicModel.forward_common_atomic (and any JAX atomic model
signatures in this module referenced from base_model.py) to accept charge_spin:
jnp.ndarray | None = None and pass it through to the internal
forward_common_atomic call so the wrapper's charge_spin kwarg is honored.
🪄 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: b79b4612-6fc3-453f-a0d5-8fb4300f0225
📒 Files selected for processing (24)
deepmd/dpmodel/descriptor/dpa3.pydeepmd/dpmodel/descriptor/make_base_descriptor.pydeepmd/entrypoints/test.pydeepmd/jax/model/base_model.pydeepmd/jax/model/dp_model.pydeepmd/jax/model/dp_zbl_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/train/training.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/dpa2.pydeepmd/pt_expt/descriptor/se_e2_a.pydeepmd/pt_expt/descriptor/se_r.pydeepmd/pt_expt/descriptor/se_t.pydeepmd/pt_expt/descriptor/se_t_tebd.pydeepmd/pt_expt/train/training.pydeepmd/utils/argcheck.pysource/tests/consistent/descriptor/common.py
🚧 Files skipped from review as they are similar to previous changes (9)
- deepmd/pt/model/descriptor/dpa1.py
- deepmd/pt/model/descriptor/se_t_tebd.py
- deepmd/pt/model/descriptor/se_t.py
- deepmd/pt/model/descriptor/dpa3.py
- deepmd/utils/argcheck.py
- deepmd/pt_expt/train/training.py
- source/tests/consistent/descriptor/common.py
- deepmd/dpmodel/descriptor/dpa3.py
- deepmd/pt/model/descriptor/se_r.py
Comment on lines
+102
to
+112
| def get_dim_chg_spin(self) -> int: | ||
| """Returns the dimension of charge_spin input (0 if not supported).""" | ||
| return 0 | ||
|
|
||
| def has_default_chg_spin(self) -> bool: | ||
| """Returns whether the descriptor has a default charge_spin value.""" | ||
| return False | ||
|
|
||
| def get_default_chg_spin(self) -> None: | ||
| """Returns the default charge_spin value, or None.""" | ||
| return None |
Contributor
There was a problem hiding this comment.
charge_spin capability is hardcoded to unsupported in a composite descriptor.
Line 102–Line 112 always return 0/False/None, which is incorrect when any child descriptor supports charge_spin. This can suppress required data plumbing for hybrid models containing DPA3-like children.
Proposed fix
def get_dim_chg_spin(self) -> int:
- """Returns the dimension of charge_spin input (0 if not supported)."""
- return 0
+ """Returns charge_spin dimension required by sub-descriptors."""
+ dims = [d.get_dim_chg_spin() for d in self.descrpt_list if d.get_dim_chg_spin() > 0]
+ if not dims:
+ return 0
+ if len(set(dims)) != 1:
+ raise ValueError(f"Inconsistent charge_spin dimensions in hybrid descriptor: {dims}")
+ return dims[0]
def has_default_chg_spin(self) -> bool:
- """Returns whether the descriptor has a default charge_spin value."""
- return False
+ """Returns whether required charge_spin-capable sub-descriptors all have defaults."""
+ req = [d for d in self.descrpt_list if d.get_dim_chg_spin() > 0]
+ return bool(req) and all(d.has_default_chg_spin() for d in req)
def get_default_chg_spin(self) -> None:
- """Returns the default charge_spin value, or None."""
- return None
+ """Returns shared default charge_spin value from required sub-descriptors, or None."""
+ req = [d for d in self.descrpt_list if d.get_dim_chg_spin() > 0]
+ if not req:
+ return None
+ vals = [d.get_default_chg_spin() for d in req]
+ if any(v is None for v in vals):
+ return None
+ first = vals[0]
+ if any(v != first for v in vals[1:]):
+ raise ValueError("Inconsistent default_chg_spin values in hybrid descriptor")
+ return first🤖 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/hybrid.py` around lines 102 - 112, The three
hardcoded accessors (get_dim_chg_spin, has_default_chg_spin,
get_default_chg_spin) must aggregate capabilities from the child descriptors
instead of returning 0/False/None; update get_dim_chg_spin to iterate your child
list (e.g., self.descriptors or whatever container the hybrid holds) and return
the sum of each child's get_dim_chg_spin(), update has_default_chg_spin to
return True if any child.has_default_chg_spin() is True, and update
get_default_chg_spin to return a combined default (e.g., concatenate
child.get_default_chg_spin() values in descriptor order) or None if no child
provides a default—handle None values when concatenating and keep ordering
consistent with how inputs are constructed.
Comment on lines
+348
to
+350
| odescriptor, gr, g2, h2, sw = descrpt( | ||
| coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin | ||
| ) |
Contributor
There was a problem hiding this comment.
Forward all supported kwargs to sub-descriptors.
Line 348 currently drops comm_dict and fparam; only charge_spin is forwarded. This silently ignores caller inputs and can break sub-descriptor behavior in hybrid mode.
Proposed fix
- odescriptor, gr, g2, h2, sw = descrpt(
- coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin
- )
+ odescriptor, gr, g2, h2, sw = descrpt(
+ coord_ext,
+ atype_ext,
+ nl,
+ mapping,
+ comm_dict=comm_dict,
+ fparam=fparam,
+ charge_spin=charge_spin,
+ )🧰 Tools
🪛 Ruff (0.15.12)
[warning] 348-348: Unpacked variable g2 is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
[warning] 348-348: Unpacked variable h2 is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
[warning] 348-348: Unpacked variable sw is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 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/hybrid.py` around lines 348 - 350, The call to
descrpt in the hybrid descriptor path only forwards charge_spin and drops other
supported keyword arguments (e.g., comm_dict, fparam), silently changing
sub-descriptor behavior; update the call site around the descrpt invocation so
it forwards the remaining kwargs—either pass comm_dict=comm_dict and
fparam=fparam (ensuring those names exist in scope) or accept and forward
**kwargs from the hybrid descriptor function into descrpt—so all supported
options reach the sub-descriptors.
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
New Features
Tests
Chores