[BUG] DPA4 native spin model raises IndexError on mixed-type frames with -1 padding

Open
  • Is something wrong?
#5,940 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
78/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Quiet
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start in deepmd/pt/model/model/sezm_native_spin_model.py at the two mask_mag call sites around lines 118-123 and 425-430, then compare the related handling in deepmd/pt/model/model/spin_model.py from #5854. Add a regression case with a mid-array -1 in source/tests/pt/model/test_spin_model_virtual_types.py and run that test; done means padded mixed-type native-spin frames no longer raise IndexError.

Written by the indexing model from the issue text.

Description

bug

SeZMNativeSpinModel builds mask_mag by indexing the per-type spin gate with the raw local atype:

https://github.com/deepmodeling/deepmd-kit/blob/777927d6805a8a4b5aeb43ba7cb46e5e0ee62ac8/deepmd/pt/model/model/sezm_native_spin_model.py#L118-L123

https://github.com/deepmodeling/deepmd-kit/blob/777927d6805a8a4b5aeb43ba7cb46e5e0ee62ac8/deepmd/pt/model/model/sezm_native_spin_model.py#L425-L430

torch.index_select rejects negative indices, unlike advanced indexing:

>>> torch.tensor([1.0, 0.0, 1.0]).index_select(0, torch.tensor([0, -1, 2]))
IndexError: index out of range in self

so any frame whose local atype contains a -1 placeholder raises instead of producing a mask.

Why this is reachable

  • DPA4/SeZM is a mixed-types descriptor, so it assumes the atom count is aligned across frames in a batch.
  • DeepmdData implements that alignment by padding with type -1: the type index map is extended with -1 in deepmd/utils/data.py, and frames are accounted with (real_type == -1).sum(axis=-1) as the ghost count.
  • model.spin.scheme: "native" is an ordinary configuration option; get_sezm_spin_model dispatches it to SeZMNativeSpinModel.

Combining those three gives a plain IndexError for a DPA4 native-spin model trained on a mixed-type system whose frames differ in atom count.

Unlike #5663 this fails loudly rather than corrupting results, so the practical effect is that the DPA4 native spin scheme cannot train on padded mixed-type data at all.

Suggested fix

The same shape as the fix applied to deepmd/pt/model/model/spin_model.py in #5854: keep index_select for export device-stability, clamp only to produce a valid index, and mask the gathered value back out.

long_atype = atype.reshape(-1).to(dtype=torch.long)
real_atom = long_atype >= 0
gathered = self.spin_mask.index_select(0, torch.clamp_min(long_atype, 0))
mask_mag = torch.where(real_atom, gathered, torch.zeros_like(gathered)).reshape(
    nf, nloc, 1
) > 0.0

Both call sites need it. A regression frame with a mid-array -1 (as in source/tests/pt/model/test_spin_model_virtual_types.py) would pin it.

Found while reviewing #5854, which fixed the same class of defect in the deepspin path; that PR deliberately left this one out of scope.

Dominant language
Python
Stars
2.1k
Forks
663
Avg merge
3d 11h
Merged PRs (30d)
25

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepmodeling/deepmd-kit

All issues in deepmodeling/deepmd-kit

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.