[BUG] Charge/spin-conditioned DPA4 breaks ragged multi-frame graph batches

[njzjz-bot]

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

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:

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

[dosubot]

Good report — the root cause and fix are clearly identified. A few related threads worth connecting:

What exists today: NeighborGraph was designed to carry a ragged n_node, and frame_id_from_n_node plus the associated segment_sum / CSR machinery are already implemented for exactly this purpose [1]. The consumer side is functionally complete and tested. The gap is that _apply_charge_spin_embedding (and the call_graph → _run_graph path that feeds it) bypasses all of that and re-derives a uniform nloc = n_out_nodes // nf, which is only correct when every frame has the same atom count [2].

Related open issues:

• [Feature Request] Neighbor-graph builders hardcode n_node = nloc, so the ragged node axis is never exercised #5938 — tracks the broader fact that all four graph builders hardcode n_node = full((nf,), nloc), so genuinely ragged batches can only be constructed manually (as in the reproducer). The fix proposed here is independent but would be the first consumer to actually require and validate a ragged graph from a non-test path [1].
• pt_expt: support correct mixed-size spin batches and ragged training #5965 — tracks the same frame_id_from_n_node-based ragged extension for the native-spin path (call_common_ragged, forward_ragged, EnergySpinLoss). The charge/spin-conditioned DPA4 issue here is a narrower instance of the same structural gap [3].
• [BUG] Multi-frame DeepPot.eval produces NaNs and corrupted outputs for frozen native-spin DPA4 graph models #5969 — reports multi-frame NaN corruption in frozen native-spin .pt2; that investigation concluded frame-offset accounting in NeighborGraph construction is sound for uniform-count batches, so those bugs are distinct from the ragged-count case described here [4].

On the suggested fix: Using frame_id_from_n_node(graph.n_node, n_total=atype.shape[0]) to index into charge_spin is the right pattern — it matches what DPA4C already does [2]. For the silent wrong-assignment case (when n_out_nodes % nf == 0 despite unequal frame counts), an assertion that graph.n_node is uniform before the nloc = n_out_nodes // nf path would make the failure loud until the ragged path is wired in. A test that runs call_graph on a ragged NeighborGraph and checks per-frame outputs against independent single-frame calls would close the coverage gap for this issue and overlap with #5938 / #5965.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.