fix(dpmodel): preserve padding when forced-sorting short neighbor lists

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)

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(