如何将2D Tensor中的值舍入到指定列表的最近值？

核心思路 #

不用循环，利用TensorFlow的广播机制和向量化操作实现，完美适配非eager模式，同时支持任意维度的输入张量。

完整代码实现 #

import tensorflow as tf

# 输入张量（支持非eager模式）
data_to_round = tf.constant([[0.3, 0.4, 2.3], [1.4, 2.2 ,55.4]])
possible_rounding_results = [1,2,3,4,5,6]

# 1. 将候选值转为张量并调整形状，适配广播计算
# 调整为(1, 1, N)形状，让2D输入张量能自动广播为(Batch, Width, N)
candidates = tf.constant(possible_rounding_results, dtype=tf.float32)
candidates_expanded = tf.expand_dims(tf.expand_dims(candidates, 0), 0)

# 2. 计算每个元素与所有候选值的绝对差
abs_diff = tf.math.abs(data_to_round[..., tf.newaxis] - candidates_expanded)

# 3. 找到每个元素对应的最小差的索引
min_indices = tf.argmin(abs_diff, axis=-1)

# 4. 根据索引提取对应的候选值，得到最终结果
rounded_data = tf.gather(candidates, min_indices)

# 验证结果（非eager模式下可使用tf.print或在Session中运行）
tf.print(rounded_data)
# 输出：[[1 1 2]
#        [1 2 5]]

关键步骤解释 #

• 广播适配：通过tf.expand_dims给候选值张量增加维度，或给输入张量用[..., tf.newaxis]扩展最后一维，让两者能进行逐元素运算，避免手动循环遍历。
• 向量化计算：所有差的绝对值计算、索引查找都是批量完成的，比循环高效得多，且完全兼容非eager的静态图模式。
• 索引提取：tf.argmin返回每个元素对应的候选值索引，再用tf.gather直接从候选列表中取出对应值，一步得到目标形状的结果。

解决你之前的痛点 #

• 无需循环：广播机制自动处理所有元素的计算，不用手动遍历1D/2D数组。
• 适配任意维度：不管输入是1D、2D还是更高维，调整扩展维度的方式即可复用代码。
• 非eager兼容：所有操作都是TensorFlow的图模式API，不依赖eager执行的动态特性，可直接在静态图中运行。

内容的提问来源于stack exchange，提问作者nhruo