deepmd.dpmodel.utils.neighbor_graph.pairs

Pairs of edges sharing a center (dst) — the edge-pair axis.

Graph-native attention uses (ordered=True, include_self=True) = the full transformer neighbor-pair square per center; 3-body angles use (ordered=False, include_self=False).

Two forms:

• compact (static_nnei=None): segment-based enumeration over the real edges only — sort edge ids by center, expand each center’s Cartesian square via cumsum offsets. Dynamic P = sum(deg**2); memory O(P) (same order as dense attention’s O(nloc * nnei**2)). The data-dependent sizes (nonzero output, tensor-repeat output) are registered as UNBACKED SymInt sizes via xp_hint_dynamic_size(), so the form traces through make_fx/torch.export and compiles under AOTI (torch >= 2.6 unbacked-symint support) — this is what makes the carry-all attention graph lower exportable to a .pt2. NumPy and JAX run it eagerly; jax.jit requires the shape-static form.

• shape-static (static_nnei set): assumes the center-major static layout (E = n_center * static_nnei, edge c * static_nnei + m belongs to center c — the layout from_dense_quartet(compact=False) emits). Pure arange/reshape arithmetic, P = n_center * static_nnei**2 with all pairs materialized and validity carried by pair_mask — no data-dependent ops at all, so it traces with only BACKED symbolic shapes (bit-parity with the dense layout; used by the dense-quartet adapter).

A global (E, E) same-center boolean is deliberately NOT used: with E ~ N * nnei it costs O(N**2 * nnei**2) memory.

Functions

center_edge_pairs(...)	Enumerate pairs of edges sharing a center.
_pairs_shape_static(...)
_pairs_compact(→ tuple[deepmd.dpmodel.array_api.Array, ...)

Module Contents

deepmd.dpmodel.utils.neighbor_graph.pairs.center_edge_pairs(dst: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, n_total: int, *, include_self: bool = True, ordered: bool = True, static_nnei: int | None = None) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array]

Enumerate pairs of edges sharing a center.

Parameters:

dstArray

(E,) int64 center of each edge (edge_index[1]).

edge_maskArray

(E,) bool, real (True) vs padding (False) edges.

n_totalint

Number of centers (bounds dst).

include_selfbool

Keep the m == n diagonal (transformer self-attention needs it).

orderedbool

Keep both (m, n) and (n, m) (attention: yes, q_m . k_n is not symmetric). False keeps only n >= m (with include_self=False: n > m — the angle set).

static_nneiint | None

None -> compact eager form. Set -> shape-static form assuming the center-major layout E = n_center * static_nnei.

Returns:

query_edgeArray

(P,) int64 edge index of the query (m).

key_edgeArray

(P,) int64 edge index of the key (n).

pair_maskArray

(P,) bool; False where either edge is padding or the pair is filtered by the include_self / ordered policy (shape-static form; the compact form drops such pairs and returns all-True).

deepmd.dpmodel.utils.neighbor_graph.pairs._pairs_shape_static(xp: Any, dev: Any, dst: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, nn: int, include_self: bool, ordered: bool) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array]

deepmd.dpmodel.utils.neighbor_graph.pairs._pairs_compact(xp: Any, dev: Any, dst: deepmd.dpmodel.array_api.Array, edge_mask: deepmd.dpmodel.array_api.Array, n_total: int, include_self: bool, ordered: bool) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array]