From 76dc26e1e33d012b1e27fa6e0deca729fa169ef8 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:26:54 +0800 Subject: [PATCH] fix(tf): initialize NVNMD map outputs Correct the MapFltNvnmd four-input contract, validate its table metadata, and zero outputs that fall outside every mapping interval. Cover deterministic allocator reuse, gradient behavior, and malformed table-gradient shapes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/map_flt_nvnmd.cc | 46 ++++++++++++++++++++-- source/tests/tf/test_nvnmd_op.py | 65 ++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/source/op/tf/map_flt_nvnmd.cc b/source/op/tf/map_flt_nvnmd.cc index 77b788e537..588e0db377 100644 --- a/source/op/tf/map_flt_nvnmd.cc +++ b/source/op/tf/map_flt_nvnmd.cc @@ -22,6 +22,8 @@ y output // //- import the library of tensorflow +#include + #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,6 +114,24 @@ class MapFltNvnmdOp : public OpKernel { auto info = t_table_info.flat().data(); auto y = t_y->flat().data(); + // Values outside every configured interval intentionally map to zero. + // Initialize first so skipped entries never 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 xx, id; diff --git a/source/tests/tf/test_nvnmd_op.py b/source/tests/tf/test_nvnmd_op.py index acf345b61d..455e0b145b 100644 --- a/source/tests/tf/test_nvnmd_op.py +++ b/source/tests/tf/test_nvnmd_op.py @@ -421,5 +421,70 @@ def test_op(self) -> None: tf.reset_default_graph() +class TestOpMapFltNvnmd(tf.test.TestCase): + """Verify the mapping op's four-input contract and range behavior.""" + + def test_out_of_range_values_map_to_zero(self) -> None: + sample_count = 4096 + x = tf.placeholder(tf.float64, [sample_count, 1]) + table = tf.constant( + [ + [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], + ], + dtype=tf.float64, + ) + table_grad = tf.constant( + [ + [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], + ], + 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), -1.0) + test_x[:5, 0] = [-1.0, 0.25, 1.25, 2.0, 3.0] + with self.cached_session() as sess: + # Reuse a nonzero, same-sized allocation so the old skipped-write + # behavior 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} + ) + + expected = np.zeros((sample_count, 1, 2)) + expected[:5, 0] = [ + [0.0, 0.0], + [10.0, 11.0], + [20.0, 21.0], + [20.0, 21.0], + [0.0, 0.0], + ] + expected_gradient = np.zeros((sample_count, 1)) + expected_gradient[:5, 0] = [0.0, 3.0, 7.0, 7.0, 0.0] + + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(actual_gradient, expected_gradient) + + def test_rejects_mismatched_table_gradient(self) -> None: + mapped = op_module.map_flt_nvnmd( + tf.zeros([1, 1], dtype=tf.float64), + tf.zeros([1, 4], dtype=tf.float64), + tf.zeros([2, 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: + with self.assertRaisesRegex( + tf.errors.InvalidArgumentError, + "table_grad shape should match table", + ): + sess.run(mapped) + + if __name__ == "__main__": unittest.main()