TensorFlow中如何沿第一轴广播张量相乘且无需perm转置？

问题原因 #

TensorFlow的广播规则是从右往左对齐维度，只有当对应维度相等或其中一方为1时才能广播。你直接执行tf.zeros([10,1]) * tf.ones([10,11,1])会失败，核心原因是维度对齐不匹配：

• 倒数第1维：两个张量都是1，符合规则；
• 倒数第2维：第一个张量是10，第二个是11，既不相等也没有一方为1，无法触发广播。

不用转置的解决方法 #

只需给形状为(10,1)的张量在中间位置增加一个维度，将其变为(10,1,1)，就能完美匹配广播规则：

• 新形状(10,1,1)与(10,11,1)从右往左对比：
  1. 倒数第1维：都是1，匹配；
  2. 倒数第2维：1和11，1可广播至11；
  3. 倒数第3维：都是10，匹配。

方法1：使用tf.expand_dims

import tensorflow as tf

# 定义原始张量
a = tf.zeros([10, 1])
b = tf.ones([10, 11, 1])

# 在axis=1位置为a增加维度，形状变为(10,1,1)
a_expanded = tf.expand_dims(a, axis=1)
result = a_expanded * b

print(result.shape)  # 输出: (10, 11, 1)

方法2：使用tf.newaxis索引扩展

这是更简洁的写法，效果与tf.expand_dims完全一致：

result = a[:, tf.newaxis, :] * b

以上两种方式都能实现沿第一个轴对应相乘、中间维度广播的需求，全程无需转置操作。

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