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.
SeZMNativeSpinModelbuildsmask_magby indexing the per-type spin gate with the raw localatype: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_selectrejects negative indices, unlike advanced indexing:so any frame whose local
atypecontains a-1placeholder raises instead of producing a mask.Why this is reachable
DeepmdDataimplements that alignment by padding with type-1: the type index map is extended with-1indeepmd/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_modeldispatches it toSeZMNativeSpinModel.Combining those three gives a plain
IndexErrorfor 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.pyin #5854: keepindex_selectfor export device-stability, clamp only to produce a valid index, and mask the gathered value back out.Both call sites need it. A regression frame with a mid-array
-1(as insource/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.