如何无循环实现TensorFlow数组指定乘法并输出单张量？

你的循环逻辑是对a的每一行（形状[1439]），与整个b（形状[1000,1439]）做逐元素乘法，最终拼接成[1000,1000,1439]的张量。以下是两种简洁的无循环实现方式：

方法1：使用tf.einsum #

TensorFlow完全兼容Numpy风格的einsum语法，直接替换np.einsum为tf.einsum即可实现：

import tensorflow as tf
from tensorflow.keras import mixed_precision

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

a = tf.random.normal(shape=[1000, 1439])
b = tf.random.normal(shape=[1000, 1439])

final_product = tf.einsum("ij, kj->ikj", a, b)
print(final_product.shape)  # 输出 (1000, 1000, 1439)

方法2：利用广播机制 #

通过给a增加一个中间维度，让TensorFlow自动完成广播匹配，无需手动复制b：

# 给a新增维度，形状变为 [1000, 1, 1439]
a_expanded = tf.expand_dims(a, axis=1)
# 广播后逐元素相乘，自动匹配维度
final_product = a_expanded * b
print(final_product.shape)  # 输出 (1000, 1000, 1439)

也可以用tf.reshape替代tf.expand_dims，写法等价：

a_reshaped = tf.reshape(a, [a.shape[0], 1, a.shape[1]])
final_product = a_reshaped * b

结果一致性验证 #

可以对比循环实现与上述方法的结果，确认完全一致：

# 原循环实现转为张量
final_product_loop = tf.stack([a[i,:] * b for i in range(b.shape[0])])
# 验证误差在可接受范围内
print(tf.reduce_all(tf.abs(final_product - final_product_loop) < 1e-6))  # 输出 True

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