From 8b30d37575cc76e289e3477f6f94cc198da68bb8 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Sun, 12 Jul 2026 01:25:01 +0800
Subject: [PATCH] fix(dpmodel): preserve short nlist padding when sorting

Fold the forced-sort regression coverage into the existing neighbor-list test matrix, including the all-padding edge case.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/dpmodel/utils/nlist.py             |  8 +++---
 source/tests/common/dpmodel/test_nlist.py | 32 +++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py
index 59e68a64a0..7af7f3b35e 100644
--- a/deepmd/dpmodel/utils/nlist.py
+++ b/deepmd/dpmodel/utils/nlist.py
@@ -296,9 +296,11 @@ def format_nlist(
     # Swapping the operands would force the SymInt comparison to run and
     # emit an `_assert_scalar` node in the exported graph.
     if extra_nlist_sort or n_nnei > nnei:
-        n_nf, n_nloc, n_nnei = nlist.shape
-        m_real_nei = nlist >= 0
-        ret = xp.where(m_real_nei, nlist, 0)
+        # Sort the padded list so forced sorting preserves the requested width
+        # when the input has fewer than ``nnei`` entries (issue #5629).
+        n_nf, n_nloc, n_nnei = ret.shape
+        m_real_nei = ret >= 0
+        ret = xp.where(m_real_nei, ret, 0)
         coord0 = xp_take_first_n(extended_coord, 1, n_nloc)
         index = xp.tile(ret.reshape(n_nf, n_nloc * n_nnei, 1), (1, 1, 3))
         coord1 = xp_take_along_axis(extended_coord, index, axis=1)
diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py
index 7f1a28e080..515c3d48c5 100644
--- a/source/tests/common/dpmodel/test_nlist.py
+++ b/source/tests/common/dpmodel/test_nlist.py
@@ -104,6 +104,38 @@ def test_nlist_st(self) -> None:
         )
         np.testing.assert_allclose(self.expected_nlist, nlist1)
 
+    def test_nlist_st_sort(self) -> None:
+        """Forced sorting must retain padding when the input is short."""
+        nlist = np.array(
+            [
+                [1, 3, -1, 2],
+                [0, -1, -1, 2],
+                [0, 1, -1, -1],
+            ],
+            dtype=np.int64,
+        ).reshape([1, self.nloc, -1])
+        nlist1 = self.md.format_nlist(
+            self.coord_ext,
+            self.atype_ext,
+            nlist,
+            extra_nlist_sort=True,
+        )
+        np.testing.assert_array_equal(self.expected_nlist, nlist1)
+
+    def test_nlist_st_sort_all_padding(self) -> None:
+        """Forced sorting must also widen an entirely padded short list."""
+        nlist = np.full((self.nf, self.nloc, 1), -1, dtype=np.int64)
+        nlist1 = self.md.format_nlist(
+            self.coord_ext,
+            self.atype_ext,
+            nlist,
+            extra_nlist_sort=True,
+        )
+        np.testing.assert_array_equal(
+            np.full_like(self.expected_nlist, -1),
+            nlist1,
+        )
+
     def test_nlist_lt(self) -> None:
         # n_nnei > nnei
         nlist = np.array(