fix(tf): honor NVNMD tensor neighbor lists (#5845)

## Summary

- parse mode-4 neighbor lists stored in NVNMD mesh tensors instead of
rebuilding them from coordinates
- validate the tensor layout, reject local or neighbor indices outside
the current frame, and own copied neighbor rows for the duration of the
op
- cover both standard and mixed-type NVNMD env-mat operators, including
malformed input

## Why existing tests missed this

Existing NVNMD entrypoint tests only exercised self-built neighbor lists
with 0-, 1-, 6-, or 7-element meshes. No NVNMD env-mat test passed a
mesh larger than 16 elements, no malformed test supplied a nonnegative
index beyond `nall`, and no caller-provided list intentionally differed
from the distance-built list. In normal inference paths those two lists
are usually identical, so rebuilding the list produced plausible output
and hid the bug.

## Validation

- `cmake --build source/build --target deepmd_op -j2`
- fresh build-tree TensorFlow op: `TestOpProdEnvMatNvnmdTensorNlist` (5
passed), including a positive out-of-range neighbor index
- `ruff format .`
- `ruff check .`
- `git diff --check`

Fixes #5652

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added support for using caller-supplied neighbor lists embedded in
mesh tensors for NVNMD environment-matrix operations (standard and
mixed-type).
* **Bug Fixes**
* Strengthened input validation for mesh-embedded neighbor data,
including truncation, out-of-range neighbor/local indices, duplicate
local indices, invalid neighbor counts, and neighbor-count overflow
beyond the provided payload.
* **Tests**
* Added new tests covering correct consumption of tensor neighbor lists
and the above failure scenarios for both operation variants.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>

Author: njzjz-bot

source/op/tf/prod_env_mat_multi_device_nvnmd.cc

@@ -16,6 +16,8 @@ date: 2021-12-6
 
 */
 
+#include <cstring>
+
 #include "coord.h"
 #include "custom_op.h"
 #include "errors.h"
@@ -112,28 +114,29 @@ static void _map_nei_info_cpu(int* nlist,
                               const bool& b_nlist_map);
 
 template <typename FPTYPE>
-static void _prepare_coord_nlist_cpu(OpKernelContext* context,
-                                     FPTYPE const** coord,
-                                     std::vector<FPTYPE>& coord_cpy,
-                                     int const** type,
-                                     std::vector<int>& type_cpy,
-                                     std::vector<int>& idx_mapping,
-                                     deepmd::InputNlist& inlist,
-                                     std::vector<int>& ilist,
-                                     std::vector<int>& numneigh,
-                                     std::vector<int*>& firstneigh,
-                                     std::vector<std::vector<int>>& jlist,
-                                     int& new_nall,
-                                     int& mem_cpy,
-                                     int& mem_nnei,
-                                     int& max_nbor_size,
-                                     const FPTYPE* box,
-                                     const int* mesh_tensor_data,
-                                     const int& nloc,
-                                     const int& nei_mode,
-                                     const float& rcut_r,
-                                     const int& max_cpy_trial,
-                                     const int& max_nnei_trial);
+static tensorflow::Status _prepare_coord_nlist_cpu(
+    FPTYPE const** coord,
+    std::vector<FPTYPE>& coord_cpy,
+    int const** type,
+    std::vector<int>& type_cpy,
+    std::vector<int>& idx_mapping,
+    deepmd::InputNlist& inlist,
+    std::vector<int>& ilist,
+    std::vector<int>& numneigh,
+    std::vector<int*>& firstneigh,
+    std::vector<std::vector<int>>& jlist,
+    int& new_nall,
+    int& mem_cpy,
+    int& mem_nnei,
+    int& max_nbor_size,
+    const FPTYPE* box,
+    const int* mesh_tensor_data,
+    const int_64 mesh_tensor_size,
+    const int& nloc,
+    const int& nei_mode,
+    const float& rcut_r,
+    const int& max_cpy_trial,
+    const int& max_nnei_trial);
 
 // instance of function
 
@@ -228,58 +231,132 @@ static void _map_nei_info_cpu(int* nlist,
                            ntypes, b_nlist_map);
 }
 
+static tensorflow::Status _prepare_mesh_nlist_cpu(
+    deepmd::InputNlist& inlist,
+    std::vector<int>& ilist,
+    std::vector<int>& numneigh,
+    std::vector<int*>& firstneigh,
+    std::vector<std::vector<int>>& jlist,
+    int& max_nbor_size,
+    const int* mesh_tensor_data,
+    const int_64 mesh_tensor_size,
+    const int nloc,
+    const int nall) {
+  // Tensor-stored lists use a 16-int header followed by ilist, numneigh,
+  // and the concatenated neighbor rows. Validate every offset before copying
+  // because this input can be constructed directly by callers.
+  const int_64 header_size = 16 + static_cast<int_64>(2) * nloc;
+  if (mesh_tensor_size < header_size) {
+    return errors::InvalidArgument("invalid mesh tensor");
+  }
+
+  const int* ilist_in = mesh_tensor_data + 16;
+  const int* numneigh_in = mesh_tensor_data + 16 + nloc;
+  const int* neighbors_in = mesh_tensor_data + header_size;
+  std::vector<int_64> neighbor_offset(nloc + 1, 0);
+  int supplied_max_nbor_size = 0;
+  for (int ii = 0; ii < nloc; ++ii) {
+    const int_64 neighbor_count = numneigh_in[ii];
+    if (neighbor_count < 0 ||
+        neighbor_offset[ii] > mesh_tensor_size - header_size - neighbor_count) {
+      return errors::InvalidArgument("invalid mesh tensor");
+    }
+    neighbor_offset[ii + 1] = neighbor_offset[ii] + neighbor_count;
+    supplied_max_nbor_size = std::max(supplied_max_nbor_size, numneigh_in[ii]);
+  }
+
+  std::vector<unsigned char> seen(static_cast<size_t>(nloc), 0);
+  for (int ii = 0; ii < nloc; ++ii) {
+    const int i_idx = ilist_in[ii];
+    if (i_idx < 0 || i_idx >= nloc || seen[static_cast<size_t>(i_idx)]) {
+      return errors::InvalidArgument("invalid mesh tensor");
+    }
+    seen[static_cast<size_t>(i_idx)] = 1;
+    for (int_64 jj = neighbor_offset[ii]; jj < neighbor_offset[ii + 1]; ++jj) {
+      // Downstream environment-matrix kernels index both coordinates and
+      // atom types with this value, so accept only live frame atoms.
+      if (neighbors_in[jj] < 0 || neighbors_in[jj] >= nall) {
+        return errors::InvalidArgument("invalid mesh tensor");
+      }
+    }
+  }
+
+  for (int ii = 0; ii < nloc; ++ii) {
+    ilist[ii] = ilist_in[ii];
+    numneigh[ii] = numneigh_in[ii];
+    jlist[ii].assign(neighbors_in + neighbor_offset[ii],
+                     neighbors_in + neighbor_offset[ii + 1]);
+    firstneigh[ii] = jlist[ii].data();
+  }
+  inlist = deepmd::InputNlist(nloc, ilist.data(), numneigh.data(),
+                              firstneigh.data());
+  max_nbor_size = std::max(max_nbor_size, supplied_max_nbor_size);
+  return tensorflow::Status();
+}
+
 template <typename FPTYPE>
-static void _prepare_coord_nlist_cpu(OpKernelContext* context,
-                                     FPTYPE const** coord,
-                                     std::vector<FPTYPE>& coord_cpy,
-                                     int const** type,
-                                     std::vector<int>& type_cpy,
-                                     std::vector<int>& idx_mapping,
-                                     deepmd::InputNlist& inlist,
-                                     std::vector<int>& ilist,
-                                     std::vector<int>& numneigh,
-                                     std::vector<int*>& firstneigh,
-                                     std::vector<std::vector<int>>& jlist,
-                                     int& new_nall,
-                                     int& mem_cpy,
-                                     int& mem_nnei,
-                                     int& max_nbor_size,
-                                     const FPTYPE* box,
-                                     const int* mesh_tensor_data,
-                                     const int& nloc,
-                                     const int& nei_mode,
-                                     const float& rcut_r,
-                                     const int& max_cpy_trial,
-                                     const int& max_nnei_trial) {
+static tensorflow::Status _prepare_coord_nlist_cpu(
+    FPTYPE const** coord,
+    std::vector<FPTYPE>& coord_cpy,
+    int const** type,
+    std::vector<int>& type_cpy,
+    std::vector<int>& idx_mapping,
+    deepmd::InputNlist& inlist,
+    std::vector<int>& ilist,
+    std::vector<int>& numneigh,
+    std::vector<int*>& firstneigh,
+    std::vector<std::vector<int>>& jlist,
+    int& new_nall,
+    int& mem_cpy,
+    int& mem_nnei,
+    int& max_nbor_size,
+    const FPTYPE* box,
+    const int* mesh_tensor_data,
+    const int_64 mesh_tensor_size,
+    const int& nloc,
+    const int& nei_mode,
+    const float& rcut_r,
+    const int& max_cpy_trial,
+    const int& max_nnei_trial) {
   inlist.inum = nloc;
-  if (nei_mode != 3) {
+  if (nei_mode != 3 && nei_mode != 4) {
     // build nlist by myself
     // normalize and copy coord
     if (nei_mode == 1) {
       int copy_ok = _norm_copy_coord_cpu(coord_cpy, type_cpy, idx_mapping,
                                          new_nall, mem_cpy, *coord, box, *type,
                                          nloc, max_cpy_trial, rcut_r);
-      OP_REQUIRES(context, copy_ok,
-                  errors::Aborted("cannot allocate mem for copied coords"));
+      if (!copy_ok) {
+        return errors::Aborted("cannot allocate mem for copied coords");
+      }
       *coord = &coord_cpy[0];
       *type = &type_cpy[0];
     }
     // build nlist
     int build_ok = _build_nlist_cpu(ilist, numneigh, firstneigh, jlist,
                                     max_nbor_size, mem_nnei, *coord, nloc,
                                     new_nall, max_nnei_trial, rcut_r);
-    OP_REQUIRES(context, build_ok,
-                errors::Aborted("cannot allocate mem for nlist"));
+    if (!build_ok) {
+      return errors::Aborted("cannot allocate mem for nlist");
+    }
     inlist.ilist = &ilist[0];
     inlist.numneigh = &numneigh[0];
     inlist.firstneigh = &firstneigh[0];
+  } else if (nei_mode == 4) {
+    tensorflow::Status status = _prepare_mesh_nlist_cpu(
+        inlist, ilist, numneigh, firstneigh, jlist, max_nbor_size,
+        mesh_tensor_data, mesh_tensor_size, nloc, new_nall);
+    if (!status.ok()) {
+      return status;
+    }
   } else {
     // copy pointers to nlist data
     memcpy(&inlist.ilist, 4 + mesh_tensor_data, sizeof(int*));
     memcpy(&inlist.numneigh, 8 + mesh_tensor_data, sizeof(int*));
     memcpy(&inlist.firstneigh, 12 + mesh_tensor_data, sizeof(int**));
     max_nbor_size = max_numneigh(inlist);
   }
+  return tensorflow::Status();
 }
 
 /*
@@ -506,11 +583,14 @@ class ProdEnvMatANvnmdQuantizeOp : public OpKernel {
         std::vector<int> type_cpy;
         int frame_nall = nall;
         // prepare coord and nlist
-        _prepare_coord_nlist_cpu<FPTYPE>(
-            context, &coord, coord_cpy, &type, type_cpy, idx_mapping, inlist,
-            ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
-            max_nbor_size, box, mesh_tensor.flat<int>().data(), nloc, nei_mode,
-            rcut_r, max_cpy_trial, max_nnei_trial);
+        OP_REQUIRES_OK(
+            context,
+            _prepare_coord_nlist_cpu<FPTYPE>(
+                &coord, coord_cpy, &type, type_cpy, idx_mapping, inlist, ilist,
+                numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
+                max_nbor_size, box, mesh_tensor.flat<int>().data(),
+                mesh_tensor.NumElements(), nloc, nei_mode, rcut_r,
+                max_cpy_trial, max_nnei_trial));
         // launch the cpu compute function
         deepmd::prod_env_mat_a_nvnmd_quantize_cpu(
             em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg,
@@ -789,11 +869,14 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel {
         std::vector<int> type_cpy;
         int frame_nall = nall;
         // prepare coord and nlist
-        _prepare_coord_nlist_cpu<FPTYPE>(
-            context, &coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist,
-            ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy, mem_nnei,
-            max_nbor_size, box, mesh_tensor.flat<int>().data(), nloc, nei_mode,
-            rcut_r, max_cpy_trial, max_nnei_trial);
+        OP_REQUIRES_OK(
+            context,
+            _prepare_coord_nlist_cpu<FPTYPE>(
+                &coord, coord_cpy, &f_type, type_cpy, idx_mapping, inlist,
+                ilist, numneigh, firstneigh, jlist, frame_nall, mem_cpy,
+                mem_nnei, max_nbor_size, box, mesh_tensor.flat<int>().data(),
+                mesh_tensor.NumElements(), nloc, nei_mode, rcut_r,
+                max_cpy_trial, max_nnei_trial));
         // launch the cpu compute function
         deepmd::prod_env_mat_a_nvnmd_quantize_cpu(
             em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg,

source/tests/tf/test_nvnmd_op.py

@@ -421,5 +421,118 @@ def test_op(self) -> None:
         tf.reset_default_graph()
 
 
+class TestOpProdEnvMatNvnmdTensorNlist(tf.test.TestCase):
+    """Verify NVNMD env-mat ops honor neighbor lists stored in mesh tensors."""
+
+    def setUp(self) -> None:
+        self.coord = np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0]])
+        self.atype = np.array([[0, 0, 0]], dtype=np.int32)
+        self.natoms = np.array([3, 3, 3], dtype=np.int32)
+        self.box = np.array([[10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0]])
+        self.sel = [2]
+        self.ndescrpt = 4 * sum(self.sel)
+        self.avg = np.zeros((1, self.ndescrpt))
+        self.std = np.ones((1, self.ndescrpt))
+
+        # All atoms are mutually within rcut, but this caller-provided list
+        # deliberately retains only one cyclic neighbor per local atom.
+        header = np.zeros(16, dtype=np.int32)
+        ilist = np.array([0, 1, 2], dtype=np.int32)
+        numneigh = np.ones(3, dtype=np.int32)
+        jlist = np.array([2, 0, 1], dtype=np.int32)
+        self.mesh = np.concatenate((header, ilist, numneigh, jlist))
+        self.expected_nlist = np.array([[2, -1, 0, -1, 1, -1]], dtype=np.int32)
+
+    def _run_op(self, mixed_types: bool, mesh: np.ndarray | None = None) -> np.ndarray:
+        op = (
+            op_module.prod_env_mat_a_mix_nvnmd_quantize
+            if mixed_types
+            else op_module.prod_env_mat_a_nvnmd_quantize
+        )
+        outputs = op(
+            tf.constant(self.coord, dtype=tf.float64),
+            tf.constant(self.atype),
+            tf.constant(self.natoms),
+            tf.constant(self.box, dtype=tf.float64),
+            tf.constant(self.mesh if mesh is None else mesh),
+            tf.constant(self.avg, dtype=tf.float64),
+            tf.constant(self.std, dtype=tf.float64),
+            rcut_a=-1.0,
+            rcut_r=3.0,
+            rcut_r_smth=0.5,
+            sel_a=self.sel,
+            sel_r=[0],
+        )
+        with self.cached_session() as sess:
+            return sess.run(outputs[3])
+
+    def test_standard_op_uses_tensor_neighbor_list(self) -> None:
+        np.testing.assert_array_equal(self._run_op(False), self.expected_nlist)
+
+    def test_mixed_type_op_uses_tensor_neighbor_list(self) -> None:
+        np.testing.assert_array_equal(self._run_op(True), self.expected_nlist)
+
+    def test_rejects_truncated_tensor_neighbor_list(self) -> None:
+        """Reject mode-4 tensors before reading an incomplete list header."""
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, self.mesh[:17])
+
+    def test_rejects_out_of_range_tensor_neighbor_index(self) -> None:
+        """Reject neighbors that would index beyond this frame's atoms."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[-1] = self.natoms[1]
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, invalid_mesh)
+
+    def test_rejects_duplicate_tensor_local_index(self) -> None:
+        """Each local atom must occur exactly once in the caller ilist."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[17] = invalid_mesh[16]
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, invalid_mesh)
+
+    def test_rejects_out_of_range_tensor_local_index(self) -> None:
+        """Reject ilist rows that do not identify a live local atom."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[16] = self.natoms[0]
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, invalid_mesh)
+
+    def test_rejects_negative_tensor_neighbor_count(self) -> None:
+        """A negative row length must never reach neighbor-row copying."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[16 + self.natoms[0]] = -1
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, invalid_mesh)
+
+    def test_rejects_tensor_neighbor_count_beyond_payload(self) -> None:
+        """Declared row lengths must fit in the remaining mesh payload."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[16 + self.natoms[0]] += 1
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(False, invalid_mesh)
+
+    def test_mixed_type_op_rejects_invalid_tensor_neighbor_list(self) -> None:
+        """The mixed-type kernel must use the same validated tensor path."""
+        invalid_mesh = self.mesh.copy()
+        invalid_mesh[-1] = self.natoms[1]
+        with self.assertRaisesRegex(
+            tf.errors.InvalidArgumentError, "invalid mesh tensor"
+        ):
+            self._run_op(True, invalid_mesh)
+
+
 if __name__ == "__main__":
     unittest.main()