feat(dpmodel): segment_max + numerically-stable mask-aware segment_softmax

Built on the existing xp_maximum_at (no new array_api helper needed).
Part of NeighborGraph PR-D (graph-native attention).

Author: Han Wang

deepmd/dpmodel/utils/neighbor_graph/__init__.py

@@ -33,7 +33,9 @@
     pad_and_guard_edges,
 )
 from .segment import (
+    segment_max,
     segment_mean,
+    segment_softmax,
     segment_sum,
 )
 
@@ -49,6 +51,8 @@
     "neighbor_graph_from_ijs",
     "node_validity_mask",
     "pad_and_guard_edges",
+    "segment_max",
     "segment_mean",
+    "segment_softmax",
     "segment_sum",
 ]

deepmd/dpmodel/utils/neighbor_graph/segment.py

@@ -9,6 +9,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_add_at,
+    xp_maximum_at,
 )
 
 
@@ -37,3 +38,52 @@ def segment_mean(data: Array, segment_ids: Array, num_segments: int) -> Array:
     # broadcast counts over the trailing dims of summed
     shape = (num_segments,) + (1,) * (summed.ndim - 1)
     return summed / xp.reshape(safe, shape)
+
+
+def segment_max(data: Array, segment_ids: Array, num_segments: int) -> Array:
+    """out[s] = max of data[i] over i with segment_ids[i] == s.
+
+    Shape ``(num_segments, *data.shape[1:])``; empty segments are ``-inf``
+    (neutral element — callers guard with masks before consuming them).
+    """
+    xp = array_api_compat.array_namespace(data)
+    out = xp.full(
+        (num_segments, *tuple(data.shape[1:])),
+        -xp.inf,
+        dtype=data.dtype,
+        device=array_api_compat.device(data),
+    )
+    return xp_maximum_at(out, segment_ids, data)
+
+
+def segment_softmax(
+    data: Array,
+    segment_ids: Array,
+    num_segments: int,
+    mask: Array | None = None,
+) -> Array:
+    """Softmax over entries sharing a segment id, numerically stable.
+
+    Mirrors the dense ``np_softmax`` max-subtraction trick with a PER-SEGMENT
+    max. ``mask`` (bool, per entry) removes masked entries from the softmax
+    entirely (zero weight AND excluded from the denominator). Empty or
+    fully-masked segments produce all-zero weights (no NaN).
+    """
+    xp = array_api_compat.array_namespace(data)
+    if mask is not None:
+        # keep masked entries out of the per-segment max: send them to -inf
+        neg = xp.full_like(data, -xp.inf)
+        data_for_max = xp.where(mask, data, neg)
+    else:
+        data_for_max = data
+    seg_max = segment_max(data_for_max, segment_ids, num_segments)
+    # guard -inf (empty / fully-masked segments) so gather doesn't yield inf-inf
+    seg_max = xp.where(xp.isinf(seg_max), xp.zeros_like(seg_max), seg_max)
+    shifted = data - xp.take(seg_max, segment_ids, axis=0)
+    ex = xp.exp(shifted)
+    if mask is not None:
+        ex = ex * xp.astype(mask, ex.dtype)
+    denom = segment_sum(ex, segment_ids, num_segments)
+    denom_e = xp.take(denom, segment_ids, axis=0)
+    safe = xp.where(denom_e > 0, denom_e, xp.ones_like(denom_e))
+    return ex / safe

source/tests/common/dpmodel/test_segment_softmax.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""segment_max / segment_softmax (NeighborGraph PR-D segment toolkit)."""
+
+import numpy as np
+
+from deepmd.dpmodel.utils.neighbor_graph import (
+    segment_max,
+    segment_softmax,
+)
+
+
+class TestSegmentMax:
+    def test_basic(self) -> None:
+        data = np.array([1.0, 5.0, 2.0, -3.0])
+        ids = np.array([0, 0, 2, 2], dtype=np.int64)
+        out = segment_max(data, ids, 3)
+        assert out[0] == 5.0
+        assert np.isneginf(out[1])  # empty segment
+        assert out[2] == 2.0
+
+    def test_trailing_dims(self) -> None:
+        data = np.array([[1.0, -2.0], [3.0, -4.0], [0.0, 9.0]])
+        ids = np.array([1, 1, 0], dtype=np.int64)
+        out = segment_max(data, ids, 2)
+        np.testing.assert_allclose(out[0], [0.0, 9.0])
+        np.testing.assert_allclose(out[1], [3.0, -2.0])
+
+    def test_torch_matches_numpy(self) -> None:
+        import torch
+
+        data = np.array([0.3, 1.2, -0.7, 2.0])
+        ids = np.array([0, 0, 1, 1], dtype=np.int64)
+        ref = segment_max(data, ids, 2)
+        out = segment_max(torch.from_numpy(data), torch.from_numpy(ids), 2)
+        np.testing.assert_allclose(out.numpy(), ref)
+
+
+class TestSegmentSoftmax:
+    def test_matches_dense(self) -> None:
+        logits = np.array([1.0, 2.0, 0.5, -1.0])
+        ids = np.array([0, 0, 0, 1], dtype=np.int64)
+        w = segment_softmax(logits, ids, 2)
+        ref0 = np.exp(np.array([1.0, 2.0, 0.5]) - 2.0)
+        ref0 = ref0 / ref0.sum()
+        np.testing.assert_allclose(w[:3], ref0, atol=1e-12)
+        np.testing.assert_allclose(w[3], 1.0, atol=1e-12)
+
+    def test_stable_large_logits(self) -> None:
+        logits = np.array([1e30, 1e30 + 1.0])
+        ids = np.array([0, 0], dtype=np.int64)
+        w = segment_softmax(logits, ids, 1)
+        assert not np.any(np.isnan(w))
+        np.testing.assert_allclose(w.sum(), 1.0, atol=1e-12)
+
+    def test_masked_entries_zero(self) -> None:
+        logits = np.array([1.0, 2.0, 3.0])
+        ids = np.array([0, 0, 0], dtype=np.int64)
+        mask = np.array([True, False, True])
+        w = segment_softmax(logits, ids, 1, mask=mask)
+        assert w[1] == 0.0
+        np.testing.assert_allclose(w.sum(), 1.0, atol=1e-12)
+        # masked entry excluded from the denominator too
+        ref = np.exp(np.array([1.0, 3.0]) - 3.0)
+        ref = ref / ref.sum()
+        np.testing.assert_allclose(w[[0, 2]], ref, atol=1e-12)
+
+    def test_all_masked_segment_is_zero_no_nan(self) -> None:
+        logits = np.array([1.0, 2.0, 5.0])
+        ids = np.array([0, 0, 1], dtype=np.int64)
+        mask = np.array([True, True, False])
+        w = segment_softmax(logits, ids, 2, mask=mask)
+        assert not np.any(np.isnan(w))
+        assert w[2] == 0.0
+
+    def test_empty_segment_no_nan(self) -> None:
+        logits = np.array([1.0, 2.0])
+        ids = np.array([0, 0], dtype=np.int64)
+        w = segment_softmax(logits, ids, 3)
+        assert not np.any(np.isnan(w))
+
+    def test_torch_matches_numpy(self) -> None:
+        import torch
+
+        logits = np.array([0.3, 1.2, -0.7, 2.0])
+        ids = np.array([0, 0, 1, 1], dtype=np.int64)
+        mask = np.array([True, True, True, False])
+        ref = segment_softmax(logits, ids, 2, mask=mask)
+        out = segment_softmax(
+            torch.from_numpy(logits),
+            torch.from_numpy(ids),
+            2,
+            mask=torch.from_numpy(mask),
+        )
+        np.testing.assert_allclose(out.numpy(), ref, atol=1e-12)