refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)

deepmd/dpmodel/array_api.py

@@ -211,6 +211,25 @@ def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
         return x
 
 
+def xp_hint_dynamic_size(x: Array) -> None:
+    """Mark a data-dependent leading dimension as a valid size for torch.export.
+
+    Under symbolic tracing (``make_fx`` / ``torch.export``) the length of a
+    data-dependent array (e.g. the output of ``nonzero`` or a tensor-``repeat``)
+    is an UNBACKED SymInt; guarding Python control flow or allocations on it
+    raises ``GuardOnDataDependentSymNode``. ``torch._check_is_size`` registers
+    the ``>= 0`` size hint that lets the tracer treat it as a proper dimension
+    (recorded as a ``sym_constrain_range_for_size`` node, preserved by AOTI).
+
+    No-op for numpy / jax / eager-torch concrete shapes — safe to call
+    unconditionally from dpmodel code (torch imported lazily, torch arrays only).
+    """
+    if array_api_compat.is_torch_array(x):
+        import torch
+
+        torch._check_is_size(x.shape[0])
+
+
 def xp_maximum_at(x: Array, indices: Array, values: Array) -> Array:
     """Segment max-assign of values into x at the specified indices.
 

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import dataclasses
 import functools
 import math
 from collections.abc import (
@@ -10,6 +9,10 @@
     Any,
 )
 
+from deepmd.dpmodel.utils.neighbor_graph import (
+    apply_pair_exclusion,
+)
+
 if TYPE_CHECKING:
     from deepmd.dpmodel.utils.neighbor_graph import (
         NeighborGraph,
@@ -33,6 +36,7 @@
 from deepmd.dpmodel.utils import (
     AtomExcludeMask,
     PairExcludeMask,
+    apply_pair_exclusion_nlist,
 )
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
@@ -294,10 +298,9 @@ def forward_common_atomic(
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
         _, nloc, _ = nlist.shape
         atype = xp_take_first_n(extended_atype, 1, nloc)
-        if self.pair_excl is not None:
-            pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype)
-            # exclude neighbors in the nlist
-            nlist = xp.where(pair_mask == 1, nlist, -1)
+        # idempotent backstop: externally-supplied nlists (C++/LAMMPS, call_lower
+        # users) bypass the in-tree builders and land here still unfiltered.
+        nlist = apply_pair_exclusion_nlist(nlist, extended_atype, self.pair_excl)
 
         ext_atom_mask = self.make_atom_mask(extended_atype)
         ret_dict = self.forward_atomic(
@@ -330,7 +333,7 @@ def forward_common_atomic_graph(
         ``self.pair_excl is not None``, an edge-keep mask is ANDed into
         ``graph.edge_mask`` before the descriptor forward, so excluded type-pairs
         contribute zero to the segment_sum. Descriptor-level ``exclude_types`` is
-        gated by ``uses_graph_lower()==False``.
+        handled inside the descriptor's ``call_graph`` (graph-native).
 
         Parameters
         ----------
@@ -356,14 +359,7 @@ def forward_common_atomic_graph(
         atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec))
         atom_mask = self.make_atom_mask(atype)  # (N,) bool
         atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype))
-        if self.pair_excl is not None:
-            keep = self.pair_excl.build_edge_exclude_mask(
-                graph.edge_index, atype_clamped
-            )
-            graph = dataclasses.replace(
-                graph,
-                edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype),
-            )
+        graph = apply_pair_exclusion(graph, atype_clamped, self.pair_excl)
         ret_dict = self.forward_atomic_graph(
             graph,
             atype_clamped,