Design proposal: unified NeighborGraph neighbor list (edges + optional angles) — one contract across backends and scenarios

[wanghan-iapcm]

Summary

This is a design proposal for a single, backend-agnostic neighbor-list contract — NeighborGraph — that replaces the divergent edge-schema work and serves every consumption scenario with one data structure and one set of operations.

Goals

1. One contract usable in training, Python inference, and C/C++ inference, each in single- and multi-frame.
2. Compile-friendly for both backends with the same contract:
   • torch / pt_expt: edge axis is a dynamic dimension Dim("nedge") (exact Σ nedge_f).
   • jax / paddle: edge axis is a static E_max capacity (jit needs static shapes).
3. Ghost-free: geometry only via per-edge displacement; force/virial scatter to local atoms.

Prior art: PR deepmodeling#5491 (pluggable NeighborList strategy + vesin), PR deepmodeling#5562 (edge schema for DPA4/SeZM — the reference this generalizes/cleans up), PR deepmodeling#5518 (edge-based force/virial).

First principles

An energy MLIP computes E = Σ_i E_i, where E_i depends only on the relative positions of atom i's neighbors within rcut. The minimal, complete, ghost-free description is the edge list:

• An edge is (src, dst) with a displacement. dst = center atom, src = neighbor. With PBC, src is a local atom and periodicity is an integer image shift S: edge_vec = r_src + S·box − r_dst. Both src, dst ∈ [0, N) are local — no ghosts.
• Geometry enters E only through edge_vec ⇒ edge_vec is the single autograd leaf.
• Force (with edge_vec = r_src − r_dst): g_e = ∂E/∂edge_vec_e; F_k = segment_sum(g_e, dst, N) − segment_sum(g_e, src, N) — scatters to local atoms only.
• Virial: global W = −Σ_e g_e ⊗ edge_vec_e; atomic = g_e ⊗ edge_vec_e split between endpoints. No absolute coordinates needed.

Consequence — the lower interface needs no coordinates and no box. Only atype (per-node, for type embedding) plus the edge list. coord/box are inputs to the builder, not to the model's lower interface.

The contract: NeighborGraph

@dataclass
class NeighborGraph:
    # per-frame sizes (ragged-native) -> flat node space N = sum(n_node)
    n_node: Array      # (nf,) nodes per frame (single-rank: local atoms; multi-rank: local+halo); rectangular => [nloc]*nf

    # 2-body edge list, flat over [0, N)
    edge_index: Array  # (2, E) int  [src, dst]; src = neighbor, dst = center; both in [0, N)
    edge_vec:   Array  # (E, 3) float  r_src - r_dst (neighbor - center); the ONLY geometry & autograd leaf
    edge_mask:  Array  # (E,) bool  real (1) vs padding (0) -- ALWAYS respected

    # multi-rank only: owned-vs-halo split (single-rank => None = all local)
    n_local: Array | None = None  # (nf,) owned atoms per frame; owned = first n_local[f] of frame f

    # OPTIONAL 3-body (angle) topology -- present iff descriptor needs angles (dpa3, se_t)
    angle_index: Array | None = None  # (2, A) [edge_a, edge_b] sharing a center (dst); into [0, E)
    angle_mask:  Array | None = None  # (A,) real vs padding on the flat angle axis

Relative to the current edge schema, the following are removed or relocated:

• coord removed — geometry is edge_vec; output shape / num_segments come from n_node.
• nall scalar removed — scatter domain = node domain = N (from n_node).
• edge_scatter_index removed — it was only needed for the hybrid "local-mapped message + extended scatter" choice. A clean design uses one index space: single-rank (i,j,S) local; multi-rank extended for both message and scatter — they coincide either way.
• atype relocated to a primary model input (gathered over [0, N)), not part of the nlist.

Derived for free: frame_id = repeat(arange(nf), n_node) (node→frame map, for per-frame fparam gather and per-frame energy reduction). n_node does double duty (sizes + frame_id).

Ghost completeness. Single-rank (incl. LAMMPS single-rank, Python PBC) is ghost-free — the builder converts host ghosts → local (i,j,S) (map ghost→owner, recover shift), so N = nloc, all nodes local, n_local=None, the core fields are complete. Multi-rank MPI cannot eliminate ghosts (cross-rank owners): the node space is extended N = nall (local+halo), and n_local is required — energy reduces over owned only, message aggregation runs over all N. Multi-rank therefore needs n_local plus a comm_dict halo layer (as in today's DPA2/DPA3); it is out of scope for the core contract.

Angle extension (3-body)

DPA3's repflows has a genuine three-body representation a^{ijk} (its own a_rcut < rcut, a_sel, update_angle; a use_dynamic_sel ragged mode vs a static nf*nloc*a_sel*a_sel dense capacity). se_t / se_t_tebd are the same structure. An angle is a pair of edges sharing a center.

• Optional angle_index (2, A) + angle_mask (A,); angle_index points into edges [0, E) (reuses edge_vec; the center is dst(edge_a) == dst(edge_b), derivable). Flat ragged angle axis A (torch dynamic / jax A_max + mask — tighter than the dense a_sel²).
• edge_vec stays the only geometry leaf — an angle's geometry (cos θ_ijk, …) is a derived function of its two referenced edge vectors, so grad(E, edge_vec) still captures angle coordinate-dependence. Force/virial are unchanged; the angle list adds topology only, no new leaf, no new backward pass.
• Angle aggregation reuses the same scatter primitive (angle→edge grouped by edge_a; angle→node by the shared center).
• Layout-driven presence: the builder emits the angle list only when a descriptor declares it needs angles; otherwise angle_index=None and edge-only descriptors ignore it.

This is the k-body hierarchy: 2-body = edges (node pairs), 3-body = angles (edge pairs sharing a node), all sharing one edge_vec geometry leaf.

Length policy: GraphLayout (the only torch/jax difference)

@dataclass
class GraphLayout:
    edge_capacity:  int | None = None   # E_max;  None -> dynamic edge axis (torch)
    angle_capacity: int | None = None   # A_max;  None -> dynamic angle axis (torch); jax pads to A_max
    node_capacity:  int | None = None   # N_max;  ragged only (rectangular: derived = nf*nloc)
    frame_capacity: int | None = None   # nf_max; ragged only (rectangular: = nf)
    min_edges: int = 2     # torch-only guard so inductor never specializes E in {0,1}

• Compact layout (all axes): [ real entries, frame-major contiguous | padding suffix ]; masks are a prefix threshold arange < real.
• torch: all None, exact counts + guard; export Dim("nedge", min=min_edges).
• jax rectangular: edge_capacity = E_max (+ angle_capacity); node/frame None (derived). jax ragged: all set.
• Overflow (Σnedge > E_max, etc.) → caller grows the capacity and recompiles outside jit (as jax-md's did_buffer_overflow), or the data loader pre-scans / bin-packs.

The data structure, the consuming code, and the math are identical across all scenarios; only the edge-axis length policy varies, and edge_mask absorbs it. This matches the convergent approach across jax-md / jraph / mace-jax / nequip-jax (flat sender/receiver list + capacity/E_max + segment_sum).

Two primitives (the real new dpmodel work)

1. edge_segment_sum(values, index, n_segments) — a backend-dispatched scatter (jax.ops.segment_sum, torch index_add_/scatter_add_, numpy np.add.at, paddle segment_sum). Padding handled by values * edge_mask before the scatter. Used for dst-aggregation, angle aggregation, per-frame energy reduction, and the force/virial scatter. This is the one non-array-API primitive the whole jax MLIP ecosystem pays for.
2. edge_energy_deriv(E, edge_vec, src, dst, edge_mask, ...) — model-level force/virial. One grad(E, edge_vec) = g_e; force, atomic virial, and global virial all fall out of that single backward. (atom_virial is therefore always computed — it is the precursor to the global virial, not an add-on, so there is no do_atomic_virial gate in the edge path.)

Lower interface — a new method, legacy forward_common_lower untouched

The edge-based lower gets a distinct name, forward_common_lower_graph; the legacy quartet forward_common_lower(extended_coord, extended_atype, nlist, mapping) is left unchanged for non-edge descriptors. Dispatch is by model type; .pt2 metadata records lower_input_kind ∈ {nlist, graph}. A neighbor_graph_from_extended adapter bridges a host-supplied quartet into the graph lower.

forward_common_lower_graph(
    atype,            # (N,)
    n_node,           # (nf,)
    edge_index,       # (2, E)
    edge_vec,         # (E, 3)   the only geometry (autograd leaf)
    edge_mask,        # (E,)
    angle_index=None, # (2, A)   optional 3-body
    angle_mask=None,  # (A,)
    n_local=None,     # (nf,)    multi-rank owned-vs-halo split
    fparam=None,      # (nf, ndf)
    aparam=None,      # (N, nda)
    comm_dict=None,   # multi-rank halo exchange
)
# outputs:
#   atom_energy (N, 1);  energy (nf, 1) = segment_sum(atom_energy, frame_id, nf) over owned
#   force       (N, 3);  atom_virial (N, 3, 3);  virial (nf, 3, 3) = reduction of atom_virial

Node tensors flatten from (nf, nloc, …) to (N, …) + n_node; edge tensors are already flat; per-frame tensors stay (nf, …). The public forward(coord, atype, box) is unchanged — it builds the graph internally and ravels/unravels at the I/O boundary.

Descriptor portability — all families port

Every deepmd descriptor is translation-invariant and depends on the environment only through r_ij = edge_vec, which the edge list carries completely; anything per-center is recoverable by grouping edges by dst. Three computation-pattern classes:

1. Factorizable two-body / message-passing — se_e2_a, se_r, dpa2, dpa3 (edge channel), DPA4/SeZM. Each edge contributes independently; aggregate with segment_sum. Even se_e2_a's bilinear D = (GᵀR)(RᵀG) factorizes (T_i = segment_sum(R_ij ⊗ G_ij), D_i = T_iᵀT_i). Trivial.
2. Attention over the neighbor set — dpa1, se_atten_v2. Per-center segment_softmax(logits, dst) (as in jraph attention).
3. Three-body / angular — se_t, se_t_tebd, the dpa3 angle channel. Served first-class by the optional angle list (built once, consumed via edge_segment_sum), not ad-hoc triplet enumeration inside each descriptor.

The port can be incremental — the legacy quartet lower coexists — and once all families are on the graph form, the quartet lower can be retired.

Decisions baked in

1. edge_vec sign = r_src − r_dst (neighbor − center; matches the existing env_mat convention).
2. edge_scatter_index removed (hybrid-only artifact).
3. nall scalar removed (scatter domain = N).
4. coord / box removed from the lower interface.
5. atom_virial always computed, no do_atomic_virial gate.
6. Retire the descriptor-internal padded edge cache in favor of NeighborGraph + edge_segment_sum (commit to a segment_sum shim in dpmodel).
7. Ragged is the native implementation; rectangular is the n_node = [nloc]*nf special case.

PR splitting and the concrete implementation plan are intentionally out of scope for this proposal — the goal here is to agree on the contract and the architecture. Feedback welcome.

[iProzd]

1. One point that I think needs to be nailed down is the semantics of sel.

The proposal describes the graph as carrying all neighbors within rcut, which is the physically complete local environment. However, in existing DeePMD descriptors, sel is not only a capacity hint; it is part of the model input contract and can affect the actual environment consumed by the descriptor, including per-type truncation and ordering behavior.

For newly trained graph-native models, using all neighbors within rcut may be the cleanest design. But for existing models or for parity with the legacy quartet lower, we need to clarify whether the graph builder should preserve the legacy sel semantics: per-type selection, truncation, sorting, and padding/masking. Otherwise, switching a model from the nlist lower to the graph lower could change numerical outputs even when the physical coordinates are unchanged.

A possible distinction might be:

• graph-native descriptors may declare sel=None / all-neighbors mode;
• legacy-compatible graph adapters must preserve the descriptor’s existing sel contract exactly.

I think this distinction should be explicit in the contract.

2. The single-rank story is very clean, but I think the multi-rank contract needs more detail before we can judge whether the same NeighborGraph is sufficient.

The proposal says that in multi-rank MPI, N = nall includes local + halo nodes, n_local marks the owned prefix, energy reduces over owned atoms only, and message aggregation runs over all N. That sounds reasonable for the node space, but several details remain important:

• If force is scattered onto halo nodes, how is the halo contribution communicated back to the owning rank?
• For multi-layer message-passing descriptors, is a similar comm_dict needed to do halo exchange between layers?
• What exactly is the boundary between the core NeighborGraph contract and the MPI/communication layer?

I agree that multi-rank can be out of scope for the first implementation, but the contract should probably state the invariants clearly enough that a later MPI implementation will not require changing the graph fields.

3. I agree with the high-level statement that DeePMD descriptors are translation-invariant and can in principle be expressed through edge_vec plus grouping by dst. However, I think the proposal may be a bit too optimistic when saying that all descriptor families can be ported straightforwardly.

For new graph-native implementations, the decomposition looks convincing. But for legacy parity, there are many details that are part of the effective descriptor contract:

• per-type sel truncation and padding;
• neighbor sorting/order assumptions;
• mixed-type vs type-split neighbor layout;
• excluded type pairs;
• virtual atoms / spin paths;
• denoising paths;
• ZBL or other pair-potential bridges;
• fparam / aparam shape and gather semantics;
• atomic output communication and masking behavior.

Even if the graph representation is mathematically complete, preserving old model outputs may require descriptor-specific adapters or compatibility modes. It would be useful to distinguish “graph-native descriptor design” from “legacy descriptor parity mode”, and to specify which level of numerical equivalence is expected during migration.

4. The GraphLayout capacity idea for JAX/Paddle makes sense, but I think the capacity policy needs to be more concrete.

Using static E_max / A_max with masks is the standard JAX-friendly approach. The main practical question is how those capacities are chosen and managed:

• Are E_max and A_max user-provided, inferred from training data, or grown automatically?
• If overflow happens during training, does the caller recompile immediately, skip the batch, or restart with a larger capacity?
• How do we avoid frequent recompilation when neighbor counts fluctuate?
• Should the data loader pre-scan / bucket / bin-pack frames by graph size?
• Do we expose overflow information explicitly, similar to did_buffer_overflow in jax-md?
• What is the expected padding overhead for dense systems or angular descriptors where A_max can grow quickly?

I think the contract is good, but without a concrete capacity-management story, the JAX/Paddle path may be hard to use reliably in production training.

[wanghan-iapcm]

Thanks — these four are the right things to pin down. The key one is sel, which also resolves much of the parity question.

1. sel — retained as a normalization factor only, never a cutoff.

sel is kept, but only as the normalization constant (the 1/nnei-style denominator inside descriptors), not as a neighbor-count cap. The graph carries all neighbors within rcut; the descriptor still divides by the configured sel. This gives a clean parity statement: when sel is large enough that the legacy descriptor never actually truncated, the graph-native descriptor reproduces the old output exactly — same neighbors, same normalization. They diverge only where the old sel was binding (truncating real neighbors), which was a modeling approximation rather than a contract worth preserving.

On your graph-native-vs-legacy-compat distinction: rather than a legacy adapter that re-implements per-type sel, the graph lower targets mixed-type descriptors only. sel collapses to a scalar normalization (sum(sel)); no per-type split, truncation, or padding slots. Neighbor ordering is irrelevant — every aggregation is a permutation-invariant segment_* reduction, so the sorted/padded (nnei) order was a dense-layout artifact, not semantics. Type-split (mixed_types=False) descriptors are not migrated; they keep the legacy quartet lower.

2. Multi-rank — comm_dict is unchanged; the graph fields don't change for MPI.

comm_dict operates on node (extended-atom) feature/force tensors over [0, N) — the forward halo-feature exchange and reverse force accumulation that today's DPA2/DPA3 border_op already does. It is orthogonal to how neighbors are represented, so it stays unchanged; the graph form only changes the edge representation, not the node halo layer. The one requirement is that the graph use the local-prefix node ordering ([0, n_local) owned, [n_local, N) halo) that comm_dict already assumes.

Concretely:

• Force on halo nodes → reverse-accumulated to owners through comm_dict (unchanged path). I'll reword "scatters to local atoms only": single-rank N = nloc so it is local; multi-rank scatters to the extended node domain [0, N), halo contributions reverse-comm to owners, and energy/virial reduce over owned ([0, n_local)) only.
• Multi-layer message passing → the same comm_dict does the per-layer halo refresh, as now.
• Boundary → NeighborGraph owns edges/angles + the n_local owned-prefix invariant; the communication layer (comm_dict) owns halo exchange. A later MPI implementation just populates n_local and supplies comm_dict — no graph field changes, which I'll state as an explicit invariant.

3. Legacy parity — agreed it is not uniformly trivial; here is what is retired vs orthogonal.

With the mixed-type + sel-as-normalization decisions above, your list splits cleanly:

• Retired (mixed-type-only): per-type sel truncation/padding, neighbor sorting/order assumptions, type-split layout. Excluded type pairs → the builder simply omits those edges.
• Orthogonal / model-level, not nlist concerns: virtual atoms / spin, denoising paths, ZBL/pair bridges, fparam/aparam (frame/node-level, gathered by frame_id), atomic-output communication — these live at model assembly and are unaffected by the graph contract. (Spin in particular is simpler in graph form: the doubled real+virtual atom set is just 2N nodes built directly, which removes the dense extend_nlist index surgery entirely; this is already demonstrated by the SeZM spin model on the edge path.)

So I would frame it as graph-native design (clean) vs out-of-scope (type-split → stays on the quartet lower), rather than a separate "legacy parity mode." Expected equivalence: exact for mixed-type descriptors in the large-sel regime.

4. Capacity policy — concrete.

• E_max/A_max/N_max: user-configurable, with a data-loader pre-scan/bin-pack to set tight per-bucket capacities for training (recommended), and auto-grow on overflow for inference.
• Training overflow → recompile with a larger capacity (not skip the batch), with a headroom factor (~1.25x) and bucketing to keep recompiles rare.
• Overflow is exposed explicitly (a did_buffer_overflow-style flag), as in jax-md.
• Angular A_max (~ sum degree^2) is the main overhead risk; mitigated by a_rcut < rcut and bin-packing, with padding fraction reported.
• The torch path is fully dynamic (no capacities) — this is a jax/paddle-only concern.

I'll fold the sel-normalization semantics, the multi-rank/comm_dict invariants, the retired-vs-orthogonal split, and the capacity policy into the spec.

[njzjz-bot]

Thanks for writing this up. I like the overall direction: a flat NeighborGraph with edge_index, edge_vec, and masks is a clean lower contract, and using edge_vec as the geometry/autograd boundary should reduce a lot of backend-specific divergence.

A few contract-level details look important to clarify before this becomes the basis for PR splitting.

edge_vec is the single autograd leaf

Consequence — the lower interface needs no coordinates and no box. Only atype ... plus the edge list.

This is the right lower-interface direction, but I would state the ownership boundary very explicitly: coord/box are removed only from the graph lower, not from the public model interface or graph builder. The builder must still own all PBC, ghost-to-owner, shift recovery, sorting/selection, and capacity overflow semantics.

edge_mask: Array # (E,) bool real (1) vs padding (0) -- ALWAYS respected

node_capacity: int | None = None

If node_capacity is part of GraphLayout, I think the contract also needs explicit node padding semantics, likely a node_mask or an equivalent invariant derived from n_node. Otherwise atype, aparam, atom_energy, force, and reductions over padded N_max are under-specified for jax ragged layouts. Edge and angle masks are not enough once nodes can be padded.

Two primitives (the real new dpmodel work)

1. edge_segment_sum(...)
2. edge_energy_deriv(...)

For the attention descriptors, segment_sum alone is not sufficient. We probably need a small segment-reduction abstraction, including at least stable segment_max / segment_softmax / masked reduce semantics, or the proposal should explicitly say where those live. Otherwise the later claim about dpa1 / se_atten_v2 portability is true architecturally but not complete as an implementation contract.

Multi-rank MPI cannot eliminate ghosts ... the node space is extended N = nall (local+halo), and n_local is required

Force ... scatters to local atoms only.

This needs careful wording. In multi-rank mode, force scatter inside the graph lower naturally lands on the extended local+halo node domain. To return owned forces, halo contributions must be reverse-accumulated to owners through the communication layer. Energy/virial reduction should be over owned centers only. I would avoid saying “scatters to local atoms only” without distinguishing single-rank local nodes from multi-rank extended nodes.

Virial: global W = −Σ_e g_e ⊗ edge_vec_e; atomic = g_e ⊗ edge_vec_e split between endpoints.

Global virial from edge gradients is fine, but atomic virial is convention-dependent. “Split between endpoints” may not match the existing DeepMD atomic virial semantics. I would make this an explicit compatibility decision rather than baking it into the new contract implicitly. Also, requiring atom_virial to be always computed may have inference/training cost; it may be better to define it as derivable from the edge path while still allowing the caller to request or skip materialization.

OPTIONAL 3-body (angle) topology
angle_index: Array | None = None

The angle abstraction is good, especially because it keeps edge_vec as the only geometry leaf. But the builder policy needs to be part of the contract: a_rcut, a_sel, type filters, whether pairs are ordered, how self-pairs are excluded, and how A ~ sum degree^2 memory/capacity overflow is handled. Otherwise different descriptors/backends may build incompatible angle graphs under the same field names.

Descriptor portability — all families port

I agree with the long-term claim, but I would phrase this more as a migration checklist. Existing descriptors have legacy assumptions around type-wise sel, neighbor ordering, padding slots, excluded types, and fixed neighbor shapes. The edge graph can represent the needed information, but the proposal should identify which assumptions are preserved by the builder and which are intentionally retired.

Finally, since the stated goal includes C/C++ inference, the ABI details should be specified early: index dtype, mask dtype, memory layout/contiguity, whether padded indices must be in-range, ownership of buffers, and whether edge_index is always [src, dst] in row-major (2, E) layout across backends.

Overall: +1 on the architecture. I would just prefer to make these invariants explicit in the NeighborGraph contract before using it as the foundation for implementation PRs.

— Authored by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5)

[wanghan-iapcm]

Thanks — agreed these invariants should be explicit in the contract before PR-splitting. Folding all of them in; point-by-point.

Ownership boundary (coord/box). Agreed, stated explicitly: coord/box are removed only from the graph lower. The public forward and the builder still own all of it — PBC, ghost→owner mapping, shift recovery, edge construction/rcut filtering, and capacity/overflow. (With sel redefined as normalization-only and mixed-type-only — see the iProzd thread — there is no per-type selection to own; the builder does rcut filtering + (i,j,S) construction, not per-type truncation/sorting.)

Node padding mask. Correct — once node_capacity (N_max) pads the node axis, edge/angle masks aren't enough. Added explicit node-padding semantics: the node mask is derivable from n_node + node_capacity as a prefix threshold (compact layout = real nodes first, padding suffix), and atype/aparam/atom_energy/force plus every N_max reduction must respect it. (Padded nodes occur only in jax ragged training, which is single-rank, so they do not co-occur with multi-rank halo.)

Segment-reduction toolkit, not just segment_sum. Right. Primitive (a) is generalized from edge_segment_sum to a small mask-aware, backend-dispatched segment-reduction module: segment_sum / segment_max / segment_mean / numerically-stable segment_softmax. That is where the attention path (dpa1/se_atten_v2) lives — the portability claim for those is now backed by a concrete primitive.

Multi-rank force wording. Agreed, reworded (also in the iProzd thread): single-rank N = nloc so the scatter is local; multi-rank scatters to the extended node domain [0, N), and halo contributions are reverse-accumulated to owners through the communication layer (comm_dict/border_op, unchanged); energy/virial reduce over owned centers ([0, n_local)) only. I will avoid "scatters to local atoms only".

Atomic virial — checked against the existing convention; it is exact parity, not a convention change. You were right to flag it, and digging in changed the answer:

• The global virial (W = -Σ_e g_e ⊗ edge_vec_e) is convention-free and always-on — sum of the per-edge outer products from the single grad(E, edge_vec), no extra backward.
• For the per-atom virial, I checked TF (prod_virial_a/r), pt-legacy (task_deriv_one + atomic_virial_corr), and the refactor(pt/dpa4): use edge schema for dpa4 and ignore sel deepmodeling/deepmd-kit#5562 edge path. The TF and pt-legacy conventions are identical — UT-proven by source/tests/consistent/model/test_ener.py, which compares atom_virial TF-vs-PT for se_e2_a/dpa1/dos/dipole/polar/zbl — and that canonical convention is the full per-edge bond virial attributed to the neighbor (src), which is displacement-based (so reproducible in the coord-free graph lower). So the graph design attributes the full -g_e ⊗ edge_vec to src, giving exact per-atom parity with the canonical TF==PT convention — not a new convention. (Note: refactor(pt/dpa4): use edge schema for dpa4 and ignore sel deepmodeling/deepmd-kit#5562's edge path currently does a 50/50 endpoint split, which is the outlier and should be corrected to full-to-src.) The atom-level N×9 materialization can still be optionally gated (zero extra backward either way); global virial is always-on.

Angle builder policy in the contract. Agreed. Pinned: a_rcut (< rcut), a_sel (normalization, consistent with the edge sel decision), type filters, unordered edge pairs, self-pair exclusion, and A ~ sum(degree^2) capacity/overflow (angle_capacity = A_max, mitigated by a_rcut < rcut + bin-packing). So all descriptors/backends build compatible angle graphs under the same field names.

Portability → migration checklist. Reframed into three buckets: graph-native (factorizable 2-body/MP, attention via segment_softmax, 3-body angles); retired (per-type sel, neighbor sorting/order, type-split layout, padding slots — the mixed-type-only decision; excluded type pairs → builder omits edges); orthogonal / model-level (spin/virtual atoms, denoising, ZBL, fparam/aparam, atomic-output comm — unaffected by the graph contract; spin is simpler in graph form). Expected equivalence: exact for mixed-type descriptors in the large-sel regime; type-split stays on the quartet lower.

C/C++ ABI, specified early. Adding an ABI section. edge_index/angle_index are int64, row-major (2, E) / (2, A) contiguous, [src, dst] — the SoA layout is deliberate: every consumer reads a whole column (index_select(edge_index[0]), segment_sum(..., edge_index[1])), so (2, E) keeps the src and dst index vectors contiguous, whereas (E, 2) would make each a strided column that forces gather/scatter copies (this matches PyG's [2, E] / jraph's separate sender/receiver arrays). edge_vec is (E, 3) (AoS) because it is read per-edge. Masks are bool/uint8; padded indices must be in-range (reserved pad node, masked); buffer ownership and contiguity fixed and identical across backends.

+1 on landing these as explicit invariants first. I'll push a contract revision with the node mask, the segment-reduction module, the full-to-src atomic virial (+ deepmodeling#5562 correction), the angle builder policy, and the ABI section.

[wanghan-iapcm]

Design updated — summary of changes from the review

Thanks to everyone's feedback above (@iProzd, @njzjz-bot). The contract has converged on several points; folding them in here so the thread reflects the current design. The top post is now stale on the items below — most importantly the atomic virial.

Substantive changes

1. Atomic virial → full-to-src (canonical convention), correcting the original "50/50 split".
The original post followed deepmodeling#5562 and split the per-edge virial 50/50 between endpoints. On checking the existing implementations, that turns out to be the outlier:

• TF (source/lib/src/prod_virial.cc) and pt-legacy (transform_output.py: task_deriv_one + atomic_virial_corr) attribute the full per-edge bond virial −g_e ⊗ edge_vec to the neighbor (src), and they are UT-proven equal by source/tests/consistent/model/test_ener.py (atom_virial compared TF-vs-PT across se_e2_a/dpa1/dos/dipole/polar/zbl).
• That canonical convention is displacement-based, hence reproducible in the coord-free graph lower.
• So the graph design attributes the full −g_e ⊗ edge_vec to src → exact per-atom parity with the canonical TF==PT convention (not a new convention). refactor(pt/dpa4): use edge schema for dpa4 and ignore sel deepmodeling/deepmd-kit#5562's 50/50 split is the lone outlier and should be corrected to full-to-src.
• The global virial is convention-free and always-on; the atom-level N×9 materialization is optionally gated (zero extra backward either way).

2. sel → normalization-only, mixed-type-only.
sel is kept purely as the normalization constant, not a neighbor-count cutoff; the graph carries all neighbors within rcut. When sel is large enough that the legacy descriptor never truncated, the graph-native descriptor reproduces the old output exactly. The graph lower targets mixed_types=True; per-type sel/truncation/sorting and type-split layout are retired (type-split models stay on the legacy quartet lower). Excluded type pairs → builder omits those edges. Neighbor ordering is irrelevant (all aggregation is permutation-invariant).

Additions

• Multi-rank / comm_dict: comm_dict is unchanged — it operates on node (extended-atom) feature/force tensors and is orthogonal to the edge representation. Multi-rank just populates n_local (owned prefix) and supplies comm_dict; no graph-field change. Force scatters to the extended node domain [0, N) with halo reverse-comm to owners; energy/virial reduce over owned ([0, n_local)).
• Capacity policy (jax/paddle): user-configurable E_max/A_max/N_max/nf_max, with data-loader bin-pack for training and auto-grow on overflow for inference; explicit did_buffer_overflow-style flag; recompile-on-overflow (not skip) with headroom + bucketing; angular A_max ~ Σ degree² mitigated by a_rcut < rcut.
• Node-padding mask: derivable from n_node + node_capacity (prefix threshold); atype/aparam/atom_energy/force and all N_max reductions respect it.
• Segment-reduction toolkit: the dpmodel scatter primitive is segment_sum / segment_max / segment_mean / numerically-stable segment_softmax (mask-aware, backend-dispatched) — segment_softmax is what the attention path (dpa1/se_atten_v2) needs.
• Angle builder policy (pinned): a_rcut (< rcut), a_sel (normalization), type filters, unordered edge pairs, self-pair exclusion, A_max overflow.
• C/C++ ABI: edge_index/angle_index are int64, row-major (2, E) / (2, A) contiguous, [src, dst] — SoA is deliberate (consumers read whole columns: index_select(edge_index[0]), segment_sum(…, edge_index[1])), matching PyG [2,E] / jraph sender-receiver arrays; edge_vec is (E, 3) (AoS, read per-edge); masks bool/uint8; padded indices in-range (reserved pad node, masked); buffers identical across backends.
• Migration checklist: graph-native (clean) / retired (per-type sel, sorting, type-split) / orthogonal-model-level (spin/virtual atoms — simpler in graph form since the doubled real+virtual set is just 2N nodes, removing the dense extend_nlist index surgery; denoising, ZBL, fparam/aparam, atomic-output comm). Expected equivalence: exact for mixed-type descriptors in the large-sel regime.

PR splitting remains out of scope for this proposal — the intent here is to converge on the contract.