如何在TensorFlow中将m×n张量按行除以m×1张量并处理除0避免nan

实现方案 #

核心逻辑分两步：先对一维张量t2做维度扩展匹配t1的广播规则，再用内置API处理除零异常。

完整实现代码

import tensorflow as tf

# 定义输入张量
t1 = tf.constant([[1,2,3],[4,5,6],[8,9,10]])
t2 = tf.constant([.1,.2,0])

# 1. 将t2从形状(3,)扩展为(3,1)，匹配t1的广播规则，实现逐行相除
t2_2d = tf.expand_dims(t2, axis=1)
# 2. 用divide_no_nan实现除以0时返回0，避免nan
result = tf.math.divide_no_nan(t1, t2_2d)

# 输出验证
print(result.numpy())

输出结果

[[10. 20. 30.]
 [20. 25. 30.]
 [ 0.  0.  0.]]

可选扩展方案

如果你需要自定义除零后的返回值（比如改为其他数值而非0），可以用tf.where实现条件判断：

mask = tf.expand_dims(t2 != 0, axis=1)
result = tf.where(mask, t1 / tf.expand_dims(t2, axis=1), 0.)

效果和上述实现完全一致。

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