fix(tf): initialize NVNMD map outputs

source/op/tf/map_flt_nvnmd.cc

@@ -22,6 +22,8 @@ y output
 //
 
 //- import the library of tensorflow
+#include <algorithm>
+
 #include "custom_op.h"
 #include "env_mat_nvnmd.h"
 
@@ -54,24 +56,42 @@ class MapFltNvnmdOp : public OpKernel {
   /// Compute the descriptor
   /// param: context
   void Compute(OpKernelContext* context) override {
-    DCHECK_EQ(3, context->num_inputs());
+    OP_REQUIRES(context, context->num_inputs() == 4,
+                deepmd::tf_compat::InvalidArgument(
+                    "MapFltNvnmd expects four input tensors"));
 
     const Tensor& t_x = context->input(0);
     const Tensor& t_table = context->input(1);
+    const Tensor& t_table_grad = context->input(2);
     const Tensor& t_table_info = context->input(3);
 
     const TensorShape& shX = t_x.shape();
     const TensorShape& shT = t_table.shape();
+    const TensorShape& shTG = t_table_grad.shape();
     const TensorShape& shI = t_table_info.shape();
 
+    OP_REQUIRES(context, shX.dims() == 2,
+                deepmd::tf_compat::InvalidArgument("Dim of x should be 2"));
+    OP_REQUIRES(context, shT.dims() == 2,
+                deepmd::tf_compat::InvalidArgument("Dim of table should be 2"));
+    OP_REQUIRES(context, shTG == shT,
+                deepmd::tf_compat::InvalidArgument(
+                    "table_grad shape should match table"));
+    OP_REQUIRES(
+        context, shI.dims() == 1,
+        deepmd::tf_compat::InvalidArgument("Dim of table_info should be 1"));
+    OP_REQUIRES(context, shT.dim_size(1) > 0 && shT.dim_size(1) % 4 == 0,
+                deepmd::tf_compat::InvalidArgument(
+                    "table width should be a positive multiple of 4"));
+    OP_REQUIRES(context, shI.dim_size(0) > 0 && shI.dim_size(0) % 5 == 0,
+                deepmd::tf_compat::InvalidArgument(
+                    "table_info length should be a positive multiple of 5"));
+
     int N = shX.dim_size(0);
     int D = shX.dim_size(1);
     int M = shT.dim_size(1) / 4;
     int S = shI.dim_size(0) / 5;
 
-    DCHECK_EQ(shX.dims(), 2);
-    DCHECK_EQ(shT.dims(), 2);
-
     /*
      * Calculate the output
      * 1.create tensor
@@ -94,55 +114,89 @@ class MapFltNvnmdOp : public OpKernel {
     auto info = t_table_info.flat<FPTYPE>().data();
     auto y = t_y->flat<FPTYPE>().data();
 
+    // Every output element is written below, but zero first so a future edit
+    // that reintroduces a skipped entry cannot expose allocator contents.
+    if (t_y->NumElements() > 0) {
+      std::fill(y, y + t_y->NumElements(), FPTYPE(0));
+    }
+
+    for (int interval = 0; interval < S; ++interval) {
+      const FPTYPE x0 = info[interval * 5 + 0];
+      const FPTYPE x1 = info[interval * 5 + 1];
+      const FPTYPE dx = info[interval * 5 + 2];
+      const int n0 = int(info[interval * 5 + 3]);
+      const int n1 = int(info[interval * 5 + 4]);
+      OP_REQUIRES(
+          context,
+          x0 <= x1 && dx > 0 && n0 >= 0 && n1 > n0 && n1 <= shT.dim_size(0),
+          deepmd::tf_compat::InvalidArgument("invalid interval in table_info"));
+    }
+
     int ss, ii, jj;
-    FPTYPE xi, x0, x1, dx;
+    FPTYPE xi, x0, dx;
     FPTYPE xx, id;
     int idx;
     int N0, N1, dN;
 
     U_Flt64_Int64 ufi;
 
     FPTYPE ytmp;
-    FPTYPE ytmp2;
-    for (ss = S - 1; ss >= 0; ss--) {
+    for (ii = 0; ii < N * D; ii++) {
+      xi = x[ii];
+      // Pick the first interval that contains xi, mirroring
+      // MapTable.mapping() in deepmd/tf/nvnmd/entrypoints/mapt.py. Intervals
+      // are ordered fine-to-coarse and share a left edge, so the first match
+      // is the most accurate one. A value outside every interval falls through
+      // to the widest (last) one and is clamped to its nearest edge below,
+      // which keeps y continuous exactly as the numpy twin does. This matters
+      // because NvnmdConfig.get_s_range() only warns when the s range exceeds
+      // the s2g table, so the closest neighbor pairs can land outside it.
+      for (ss = 0; ss < S - 1; ss++) {
+        if ((xi >= info[ss * 5 + 0]) && (xi <= info[ss * 5 + 1])) {
+          break;
+        }
+      }
       x0 = info[ss * 5 + 0];
-      x1 = info[ss * 5 + 1];
       dx = info[ss * 5 + 2];
       N0 = int(info[ss * 5 + 3]);
       N1 = int(info[ss * 5 + 4]);
       dN = N1 - N0;
-      for (ii = 0; ii < N * D; ii++) {
-        // cal idx and xx
-        xi = x[ii];
-        if ((xi < x0) || (xi > x1)) {
-          continue;
-        }
-        //
-        xx = xi - x0;
-        id = floor(xx / dx);
-        id = (id < 0) ? 0 : id;
-        id = (id >= dN) ? (dN - 1) : id;
+      // cal idx and xx
+      xx = xi - x0;
+      id = floor(xx / dx);
+      // Written as !(id >= 0) so NaN also lands on a valid table edge rather
+      // than reaching the undefined integer conversion below.
+      if (!(id >= FPTYPE(0))) {
+        // Below the table: evaluate at the left edge of the selected interval.
+        id = 0;
+        xx = FPTYPE(0);
+      } else if (id >= dN) {
+        // Above the table: evaluate at the right edge of the last row, which
+        // is what mapt.py does via idx_k = N1 - 1 with dxx_k = dx.
+        id = dN - 1;
+        xx = dx;
+      } else {
         xx -= id * dx;
-        idx = id + N0;
-        //
-        ufi.nflt = xx;
-        ufi.nint &= 0xfffffff000000000;  // 52 - 16 = 36 = 9 * 4
-        xx = ufi.nflt;
-        for (jj = 0; jj < M; jj++) {
-          FPTYPE a = table[idx * M * 4 + jj * 4 + 0];
-          FPTYPE b = table[idx * M * 4 + jj * 4 + 1];
-          FPTYPE c = table[idx * M * 4 + jj * 4 + 2];
-          FPTYPE d = table[idx * M * 4 + jj * 4 + 3];
-          mul_flt_nvnmd(ytmp, a, xx);
-          add_flt_nvnmd(ytmp, b, ytmp);
-          mul_flt_nvnmd(ytmp, ytmp, xx);
-          add_flt_nvnmd(ytmp, c, ytmp);
-          mul_flt_nvnmd(ytmp, ytmp, xx);
-          add_flt_nvnmd(ytmp, d, ytmp);
-          y[ii * M + jj] = ytmp;
-        }  // jj
-      }  // ii
-    }  // ss
+      }
+      idx = int(id) + N0;
+      //
+      ufi.nflt = xx;
+      ufi.nint &= 0xfffffff000000000;  // 52 - 16 = 36 = 9 * 4
+      xx = ufi.nflt;
+      for (jj = 0; jj < M; jj++) {
+        FPTYPE a = table[idx * M * 4 + jj * 4 + 0];
+        FPTYPE b = table[idx * M * 4 + jj * 4 + 1];
+        FPTYPE c = table[idx * M * 4 + jj * 4 + 2];
+        FPTYPE d = table[idx * M * 4 + jj * 4 + 3];
+        mul_flt_nvnmd(ytmp, a, xx);
+        add_flt_nvnmd(ytmp, b, ytmp);
+        mul_flt_nvnmd(ytmp, ytmp, xx);
+        add_flt_nvnmd(ytmp, c, ytmp);
+        mul_flt_nvnmd(ytmp, ytmp, xx);
+        add_flt_nvnmd(ytmp, d, ytmp);
+        y[ii * M + jj] = ytmp;
+      }  // jj
+    }  // ii
   }  // Compute
 };  // MapFltNvnmdOp
 

source/tests/tf/test_nvnmd_op.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import os
 import unittest
+from typing import (
+    ClassVar,
+)
 
 import numpy as np
 
@@ -421,6 +424,161 @@ def test_op(self) -> None:
         tf.reset_default_graph()
 
 
+class TestOpMapFltNvnmd(tf.test.TestCase):
+    """Verify the mapping op's input contract and out-of-range behavior."""
+
+    # Two rows of two output channels; every polynomial is the constant d, so
+    # the mapped value identifies the selected table row.
+    table_values: ClassVar[list[list[float]]] = [
+        [0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 11.0],
+        [0.0, 0.0, 0.0, 20.0, 0.0, 0.0, 0.0, 21.0],
+    ]
+    table_grad_values: ClassVar[list[list[float]]] = [
+        [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 2.0],
+        [0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 4.0],
+    ]
+
+    def test_out_of_range_values_clamp_to_the_table_edge(self) -> None:
+        """Below/above the table evaluates the first/last row, as mapt.py does."""
+        sample_count = 4096
+        x = tf.placeholder(tf.float64, [sample_count, 1])
+        table = tf.constant(self.table_values, dtype=tf.float64)
+        table_grad = tf.constant(self.table_grad_values, dtype=tf.float64)
+        table_info = tf.constant([0.0, 2.0, 1.0, 0.0, 2.0], dtype=tf.float64)
+
+        mapped = op_module.map_flt_nvnmd(x, table, table_grad, table_info)
+        gradient = tf.gradients(tf.reduce_sum(mapped), x)[0]
+
+        warm_x = np.tile([[0.25], [1.25]], (sample_count // 2, 1))
+        test_x = np.full((sample_count, 1), 0.25)
+        test_x[:6, 0] = [np.nan, -1.0, 0.25, 1.25, 2.0, 3.0]
+        with self.cached_session() as sess:
+            # Reuse a nonzero, same-sized allocation so a skipped write cannot
+            # pass merely because fresh pages happen to be zero.
+            sess.run(mapped, feed_dict={x: warm_x})
+            actual, actual_gradient = sess.run(
+                [mapped, gradient], feed_dict={x: test_x}
+            )
+
+        # NaN and -1 clamp to row 0, while 2.0 and 3.0 clamp to row 1.
+        expected = np.tile([[[10.0, 11.0]]], (sample_count, 1, 1))
+        expected[:6, 0] = [
+            [10.0, 11.0],
+            [10.0, 11.0],
+            [10.0, 11.0],
+            [20.0, 21.0],
+            [20.0, 21.0],
+            [20.0, 21.0],
+        ]
+        expected_gradient = np.full((sample_count, 1), 3.0)
+        expected_gradient[:6, 0] = [3.0, 3.0, 3.0, 7.0, 7.0, 7.0]
+
+        np.testing.assert_array_equal(actual, expected)
+        np.testing.assert_array_equal(actual_gradient, expected_gradient)
+
+    def test_first_matching_interval_wins(self) -> None:
+        """Nested fine/coarse intervals resolve like MapTable.mapping()."""
+        # Fine interval [0, 1] over row 0 only; coarse interval [0, 4] over
+        # row 1 only. 0.5 is inside both and must take the fine one.
+        table_info = tf.constant(
+            [0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 4.0, 4.0, 1.0, 2.0],
+            dtype=tf.float64,
+        )
+        x = tf.constant([[0.5], [2.0], [9.0]], dtype=tf.float64)
+        mapped = op_module.map_flt_nvnmd(
+            x,
+            tf.constant(self.table_values, dtype=tf.float64),
+            tf.constant(self.table_grad_values, dtype=tf.float64),
+            table_info,
+        )
+        with self.cached_session() as sess:
+            actual = sess.run(mapped)
+
+        np.testing.assert_array_equal(
+            actual,
+            [[[10.0, 11.0]], [[20.0, 21.0]], [[20.0, 21.0]]],
+        )
+
+    def _run_invalid(self, message: str, **overrides: tf.Tensor) -> None:
+        """Run the op with one input replaced and require it to be rejected."""
+        inputs = {
+            "x": tf.zeros([1, 1], dtype=tf.float64),
+            "table": tf.zeros([1, 4], dtype=tf.float64),
+            "table_grad": tf.zeros([1, 4], dtype=tf.float64),
+            "table_info": tf.constant([0.0, 1.0, 1.0, 0.0, 1.0], dtype=tf.float64),
+        }
+        inputs.update(overrides)
+        mapped = op_module.map_flt_nvnmd(**inputs)
+        with self.cached_session() as sess:
+            with self.assertRaisesRegex(tf.errors.InvalidArgumentError, message):
+                sess.run(mapped)
+
+    def test_accepts_the_minimal_valid_inputs(self) -> None:
+        """The shared fixture must pass, so each rejection isolates one cause."""
+        mapped = op_module.map_flt_nvnmd(
+            tf.zeros([1, 1], dtype=tf.float64),
+            tf.zeros([1, 4], dtype=tf.float64),
+            tf.zeros([1, 4], dtype=tf.float64),
+            tf.constant([0.0, 1.0, 1.0, 0.0, 1.0], dtype=tf.float64),
+        )
+        with self.cached_session() as sess:
+            np.testing.assert_array_equal(sess.run(mapped), np.zeros([1, 1, 1]))
+
+    def test_rejects_wrong_x_rank(self) -> None:
+        self._run_invalid("Dim of x should be 2", x=tf.zeros([1], dtype=tf.float64))
+
+    def test_rejects_wrong_table_rank(self) -> None:
+        self._run_invalid(
+            "Dim of table should be 2",
+            table=tf.zeros([1, 1, 4], dtype=tf.float64),
+            table_grad=tf.zeros([1, 1, 4], dtype=tf.float64),
+        )
+
+    def test_rejects_mismatched_table_gradient(self) -> None:
+        self._run_invalid(
+            "table_grad shape should match table",
+            table_grad=tf.zeros([2, 4], dtype=tf.float64),
+        )
+
+    def test_rejects_wrong_table_info_rank(self) -> None:
+        self._run_invalid(
+            "Dim of table_info should be 1",
+            table_info=tf.constant([[0.0, 1.0, 1.0, 0.0, 1.0]], dtype=tf.float64),
+        )
+
+    def test_rejects_table_width_not_multiple_of_four(self) -> None:
+        self._run_invalid(
+            "table width should be a positive multiple of 4",
+            table=tf.zeros([1, 5], dtype=tf.float64),
+            table_grad=tf.zeros([1, 5], dtype=tf.float64),
+        )
+
+    def test_rejects_table_info_length_not_multiple_of_five(self) -> None:
+        self._run_invalid(
+            "table_info length should be a positive multiple of 5",
+            table_info=tf.constant([0.0, 1.0, 1.0, 0.0], dtype=tf.float64),
+        )
+
+    def test_rejects_each_invalid_interval_field(self) -> None:
+        """Cover every sub-condition of the per-interval table_info check."""
+        # x0 x1 dx N0 N1, one broken field per case.
+        cases = {
+            "x0 > x1": [1.0, 0.0, 1.0, 0.0, 1.0],
+            "dx == 0": [0.0, 1.0, 0.0, 0.0, 1.0],
+            "dx < 0": [0.0, 1.0, -1.0, 0.0, 1.0],
+            "N0 < 0": [0.0, 1.0, 1.0, -1.0, 1.0],
+            "N1 == N0": [0.0, 1.0, 1.0, 0.0, 0.0],
+            "N1 < N0": [0.0, 1.0, 1.0, 1.0, 0.0],
+            "N1 > rows": [0.0, 1.0, 1.0, 0.0, 2.0],
+        }
+        for name, info in cases.items():
+            with self.subTest(case=name):
+                self._run_invalid(
+                    "invalid interval in table_info",
+                    table_info=tf.constant(info, dtype=tf.float64),
+                )
+
+
 class TestOpProdEnvMatNvnmdTensorNlist(tf.test.TestCase):
     """Verify NVNMD env-mat ops honor neighbor lists stored in mesh tensors."""