From 76dc26e1e33d012b1e27fa6e0deca729fa169ef8 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:26:54 +0800 Subject: [PATCH 1/3] 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() From 41edda99f74a8abe407386dcbe4cc8ef51174ebb Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 12:36:05 +0800 Subject: [PATCH 2/3] fix(tf): clamp NVNMD map inputs to the table edge Values outside every configured interval mapped to zero, which is a legitimate embedding output rather than a sentinel and diverges from MapTable.mapping() in mapt.py. Select the first containing interval per element and clamp to its nearest edge instead, keeping y continuous and matching the numpy twin. NvnmdConfig.get_s_range() only warns when the s range exceeds the s2g table, so this path is reachable. Cover every OP_REQUIRES site, including each sub-condition of the per-interval table_info check. --- source/op/tf/map_flt_nvnmd.cc | 88 +++++++++-------- source/tests/tf/test_nvnmd_op.py | 156 ++++++++++++++++++++++++------- 2 files changed, 175 insertions(+), 69 deletions(-) diff --git a/source/op/tf/map_flt_nvnmd.cc b/source/op/tf/map_flt_nvnmd.cc index 588e0db377..4274c80804 100644 --- a/source/op/tf/map_flt_nvnmd.cc +++ b/source/op/tf/map_flt_nvnmd.cc @@ -114,8 +114,8 @@ 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. + // 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)); } @@ -133,7 +133,7 @@ class MapFltNvnmdOp : public OpKernel { } int ss, ii, jj; - FPTYPE xi, x0, x1, dx; + FPTYPE xi, x0, dx; FPTYPE xx, id; int idx; int N0, N1, dN; @@ -141,46 +141,60 @@ class MapFltNvnmdOp : public OpKernel { 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); + if (id < 0) { + // Below the table: evaluate at the left edge of the first 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 diff --git a/source/tests/tf/test_nvnmd_op.py b/source/tests/tf/test_nvnmd_op.py index 455e0b145b..ad670fa26e 100644 --- a/source/tests/tf/test_nvnmd_op.py +++ b/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 @@ -422,68 +425,157 @@ def test_op(self) -> None: 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: + """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( - [ - [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 = 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), -1.0) + test_x = np.full((sample_count, 1), 0.25) 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. + # 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} ) - expected = np.zeros((sample_count, 1, 2)) + # -1 clamps to row 0, 2.0 and 3.0 clamp to row 1; nothing maps to zero. + expected = np.tile([[[10.0, 11.0]]], (sample_count, 1, 1)) expected[:5, 0] = [ - [0.0, 0.0], [10.0, 11.0], + [10.0, 11.0], + [20.0, 21.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] + expected_gradient = np.full((sample_count, 1), 3.0) + expected_gradient[:5, 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_rejects_mismatched_table_gradient(self) -> None: + 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([2, 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: - with self.assertRaisesRegex( - tf.errors.InvalidArgumentError, - "table_grad shape should match table", - ): - sess.run(mapped) + 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), + ) if __name__ == "__main__": From 17319eddfa59f72ea80d513bbd0ac16f0106a850 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 11:35:18 +0800 Subject: [PATCH 3/3] fix(tf): clamp NaN NVNMD map inputs Address the outstanding requested-change review comments. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/map_flt_nvnmd.cc | 6 ++++-- source/tests/tf/test_nvnmd_op.py | 9 +++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/source/op/tf/map_flt_nvnmd.cc b/source/op/tf/map_flt_nvnmd.cc index 4274c80804..0283401974 100644 --- a/source/op/tf/map_flt_nvnmd.cc +++ b/source/op/tf/map_flt_nvnmd.cc @@ -164,8 +164,10 @@ class MapFltNvnmdOp : public OpKernel { // cal idx and xx xx = xi - x0; id = floor(xx / dx); - if (id < 0) { - // Below the table: evaluate at the left edge of the first interval. + // 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) { diff --git a/source/tests/tf/test_nvnmd_op.py b/source/tests/tf/test_nvnmd_op.py index ad670fa26e..a77218af64 100644 --- a/source/tests/tf/test_nvnmd_op.py +++ b/source/tests/tf/test_nvnmd_op.py @@ -451,7 +451,7 @@ def test_out_of_range_values_clamp_to_the_table_edge(self) -> None: warm_x = np.tile([[0.25], [1.25]], (sample_count // 2, 1)) test_x = np.full((sample_count, 1), 0.25) - test_x[:5, 0] = [-1.0, 0.25, 1.25, 2.0, 3.0] + 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. @@ -460,9 +460,10 @@ def test_out_of_range_values_clamp_to_the_table_edge(self) -> None: [mapped, gradient], feed_dict={x: test_x} ) - # -1 clamps to row 0, 2.0 and 3.0 clamp to row 1; nothing maps to zero. + # 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[:5, 0] = [ + expected[:6, 0] = [ + [10.0, 11.0], [10.0, 11.0], [10.0, 11.0], [20.0, 21.0], @@ -470,7 +471,7 @@ def test_out_of_range_values_clamp_to_the_table_edge(self) -> None: [20.0, 21.0], ] expected_gradient = np.full((sample_count, 1), 3.0) - expected_gradient[:5, 0] = [3.0, 3.0, 7.0, 7.0, 7.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)