deepmodeling / deepmodeling/deepmd-kit
[BUG] Charge/spin-conditioned DPA4 breaks ragged multi-frame graph batches
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### Summary
Charge/spin-conditioned DPA4 assumes every frame in a flat graph batch has the same node count. It crashes for a genuinely ragged batch and can attach frame conditions to the wrong nodes for other unequal-count batches.
`NeighborGraph` defines a flat node axis with `N = sum(n_node)`, and the public `EnergyModel.forward_ragged()` accepts arbitrary per-frame `n_node` together with `(nf, 2)` `charge_spin`. This reproduces on `origin/master` at `8cfd46e37448`.
### Reproduction
```python
import torch
from deepmd.dpmodel.utils.neighbor_graph import NeighborGraph
from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4
d = DescrptDPA4(
ntypes=1,
sel=2,
rcut=3.0,
channels=4,
n_radial=2,
lmax=0,
kmax=0,
n_blocks=0,
use_env_seed=False,
random_gamma=False,
add_chg_spin_ebd=True,
precision="float64",
seed=1,
).eval()
graph = NeighborGraph(
n_node=torch.tensor([1, 2], dtype=torch.int64),
edge_index=torch.zeros((2, 2), dtype=torch.int64),
edge_vec=torch.zeros((2, 3), dtype=torch.float64),
edge_mask=torch.zeros((2,), dtype=torch.bool),
)
d.call_graph(
graph,
torch.zeros((3,), dtype=torch.int64),
charge_spin=torch.tensor([[0.0, 1.0], [1.0, 2.0]], dtype=torch.float64),
)
```
Actual result:
```text
deepmd/dpmodel/descriptor/dpa4.py:1972
RuntimeError: shape '[3, 4]' is invalid for input of size 8
```
### Cause
`call_graph()` reduces the graph to scalar `nf`; `_run_graph()` then passes `nloc=n_out_nodes // nf`; `_apply_charge_spin_embedding()` broadcasts every frame condition to that uniform width and reshapes it onto the flat node axis.
For `n_node=[1, 2]`, this creates only two condition rows for three nodes. If the total happens to be divisible by `nf`, unequal frame counts can instead silently assign conditions to the wrong nodes.
### Suggested fix
Gather the per-frame condition with `frame_id_from_n_node(graph.n_node, n_total=atype.shape[0])`, as DPA4C does, and compare a ragged batch against independent single-frame calls in descriptor- and model-level tests.
Related to the currently unreachable production ragged builders tracked in #5938, but this is a separate consumer correctness failure in the public ragged API.
---
Coding agent: Codex
Codex version: codex-cli 0.149.0
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Assessment
This issue has not been assessed yet.