deepmodeling / deepmodeling/deepmd-kit
[BUG] DPA4 native spin model raises IndexError on mixed-type frames with -1 padding
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
`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.
```python
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.
Contributor guide
Assessment
This issue has not been assessed yet.