TensorFlow中不同形状矩阵逐行element-wise乘积的张量实现方法

嘿，这个需求其实用TensorFlow的广播机制或者tf.einsum都能很优雅地实现，完全不用手动写循环，效率还高！

核心思路 #

假设矩阵M的形状是(a, d)（a行d列），矩阵N的形状是(b, d)（b行d列），我们要生成的目标张量形状是(a, b, d)，其中每个(i,j)位置的d维向量，就是M的第i行和N的第j行逐元素相乘的结果。

方法一：利用广播机制（最基础直观） #

广播是TensorFlow自动扩展张量维度以匹配运算的特性，我们只需要给M和N分别增加一个维度，让它们能“对齐”后逐元素相乘：

1. 给M的列维度前增加一个维度，把(a,d)变成(a,1,d)
2. 给N的行维度前增加一个维度，把(b,d)变成(1,b,d)
3. 直接相乘，TensorFlow会自动把两个张量广播成(a,b,d)的形状，然后逐元素计算乘积

代码示例：

import tensorflow as tf

# 定义示例矩阵
M = tf.constant([[1, 2], [3, 4], [5, 6]])  # shape: (3, 2)
N = tf.constant([[7, 8], [9, 10]])         # shape: (2, 2)

# 扩展维度
M_expanded = tf.expand_dims(M, axis=1)  # shape: (3, 1, 2)
N_expanded = tf.expand_dims(N, axis=0)  # shape: (1, 2, 2)

# 逐元素相乘，自动广播
result = M_expanded * N_expanded  # shape: (3, 2, 2)

# 验证结果：第0行第1列的元素是M[0] * N[1] = [1*9, 2*10]
print(result[0, 1].numpy())  # 输出: [ 9 20]

方法二：使用tf.einsum（更灵活的符号式写法） #

如果你熟悉爱因斯坦求和符号，tf.einsum会是更简洁的选择，它可以直接通过字符串定义张量的运算规则：

result = tf.einsum('id,jd->ijd', M, N)

这里的符号解释：

• 'id'：表示M的维度是i（行）和d（列）
• 'jd'：表示N的维度是j（行）和d（列）
• 'ijd'：表示输出张量的维度是i、j、d，对应M的第i行和N的第j行逐元素相乘的结果

两种方法得到的结果完全一致，你可以根据自己的习惯选择～

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