Skip to content

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

Description

@wanghan-iapcm

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions