如何在TensorFlow中实现指定维度的矩阵式逐元素相乘？

这个需求其实用TensorFlow的广播机制或者tf.einsum就能轻松搞定，我给你两种实用的方案：

方案一：利用广播机制（最直观） #

核心思路是给两个张量分别扩展一个维度，让它们能够通过广播自动匹配到目标形状，然后直接做逐元素相乘：

1. 给张量A在第2个维度（也就是b_len对应的位置）增加一个维度，把形状从bsz × a_len × dim变成bsz × a_len × 1 × dim
2. 给张量B在第1个维度（也就是a_len对应的位置）增加一个维度，把形状从bsz × b_len × dim变成bsz × 1 × b_len × dim
3. 直接对扩展后的张量做逐元素乘法，TensorFlow的广播会自动把两个张量扩展到bsz × a_len × b_len × dim的形状，完成每个向量对的逐元素相乘

示例代码如下：

import tensorflow as tf

# 构造示例张量（你可以替换成自己的张量）
bsz = 2
a_len = 3
b_len = 4
dim = 5

A = tf.random.normal((bsz, a_len, dim))
B = tf.random.normal((bsz, b_len, dim))

# 扩展维度
A_expanded = tf.expand_dims(A, axis=2)  # 形状变为 (2, 3, 1, 5)
B_expanded = tf.expand_dims(B, axis=1)  # 形状变为 (2, 1, 4, 5)

# 逐元素相乘得到结果
result = A_expanded * B_expanded

# 验证形状是否符合预期
print(result.shape)  # 输出 (2, 3, 4, 5)

方案二：使用tf.einsum（更简洁） #

如果你喜欢更紧凑的代码，可以用tf.einsum，它通过爱因斯坦求和符号来指定维度的映射关系，写法非常简洁：

import tensorflow as tf

# 同样用上面的示例张量
result = tf.einsum('bad,bcd->bacd', A, B)
print(result.shape)  # 同样输出 (2, 3, 4, 5)

这里的符号解释一下：

• 'bad'对应A的维度(bsz, a_len, dim)
• 'bcd'对应B的维度(bsz, b_len, dim)
• 'bacd'是我们想要的输出维度(bsz, a_len, b_len, dim)
  einsum会自动处理维度的匹配和逐元素相乘，结果和方案一完全一致。

两种方案都是高效的，广播机制底层不会复制额外数据，einsum则是TensorFlow优化过的操作，你可以根据自己的习惯选择。

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