如何利用TensorFlow/Keras广播实现特定形状张量的外积运算？

将现有两个张量：张量a的形状为(n, 1)，张量b的形状为(p, q)，其中p和q取值较小但n取值较大。希望得到它们的“外积”张量c，其形状为(n, p, q)，满足元素关系c[i, j, k] = a[i, 0] * b[j, k]。请问如何在不使用循环的前提下，利用TensorFlow/Keras的广播机制或相关函数实现该需求？

这事儿用TensorFlow的广播机制就能轻松搞定，完全不用写循环，核心就是调整一下张量的维度，让它们符合广播规则就行，我给你两种实用的实现方式：

方法一：手动扩维 + 直接相乘（推荐，内存效率更高） #

TensorFlow的广播规则允许形状匹配的张量进行元素级运算，我们只需要把张量a的维度调整成能和b匹配的样子：

• 张量a原本是(n, 1)，我们给它在最后再加一个长度为1的维度，变成(n, 1, 1)
• 张量b的形状是(p, q)，广播机制会自动把它的形状逻辑上扩展为(1, p, q)
• 两者相乘时，每个维度都会自动匹配扩展，最终得到(n, p, q)的结果

代码示例：

import tensorflow as tf

# 随便定义两个测试张量，n取1000（模拟大值），p=3，q=4
a = tf.random.normal((1000, 1))
b = tf.random.normal((3, 4))

# 给a扩展最后一个维度，两种写法都可以
a_expanded = tf.expand_dims(a, axis=-1)  # 写法1
# a_expanded = a[..., tf.newaxis]  # 写法2，更简洁

# 直接相乘，广播会自动处理维度匹配
c = a_expanded * b

print(c.shape)  # 输出 (1000, 3, 4)，完全符合需求

这个方法的优势是不需要实际复制数据，广播只是在计算时做逻辑上的维度扩展，对于n很大的场景，内存效率比显式复制高很多。

方法二：使用tf.tile显式复制（不推荐大n场景） #

如果你想显式复制a的维度来匹配b的形状，也可以用tf.tile，但这种方法会实际复制数据，当n很大时会占用更多内存，只做了解即可：

import tensorflow as tf

n, p, q = 1000, 3, 4
a = tf.random.normal((n, 1))
b = tf.random.normal((p, q))

# 先把a扩维成(n,1,1)，然后在第1、2维度分别复制p、q次
a_tiled = tf.tile(tf.expand_dims(a, axis=-1), multiples=[1, p, q])
c = a_tiled * b

print(c.shape)  # 同样输出 (1000, 3, 4)

另外补充：如果用Keras的话，这些操作完全兼容，因为Keras的张量本质就是TensorFlow张量（TensorFlow作为后端时），直接套用上面的代码就行。

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