Numpy 使用不同形状张量的tf.einsum

我有两个张量，A和B：

A = tf.Variable(initial_value=tf.random_normal_initializer()(shape=(1, 64, 4, 1, 1)))
B = tf.Variable(initial_value=tf.random_normal_initializer()(shape=(32, 1, 4, 256, 256)))

我正在尝试执行此操作：

tf.einsum('abcde,abcde->abde', A, B) # in other words, "K.sum(A * B, axis=2)"

但是当我运行代码时，我得到一个错误：

输入形状[32,1,4256256]的轴0处的预期尺寸1，但获得尺寸32[Op:Einsum]

我不明白发生了什么，因为同样的代码使用numpy工作：

A = np.ones((1, 64, 4, 1, 1))
B = np.ones((32, 1, 4, 256, 256))

np.einsum('abcde,abcde->abde', A, B) # works great

---

为什么张量在张量流上需要有相同的形状？我不希望它们具有相同的形状，因为它会消耗大量内存，numpy上的相同代码工作得非常完美。

tf.einsum（）

似乎不支持广播，而

np.einsum（）

支持广播。如果在

K.sum（A*B，axis=2）

中有一个替代方法，为什么不使用它呢？我在问题上解释过，如果我使用

K.sum（A*B，axis=2）

，乘法A*B将创建一个形状为（32，64，4，256，256）的张量，这将消耗大量内存，这并不是必须的，因为我将在下一步对轴2求和，所以最好的解决方案是直接使用einsum获得结果（32、64、256、256），并避免创建5D数组。您不需要保持大小为1的维度，

np.einsum（'bc，acde->abd'，np.squence（A），np.squence（B））

同样有效。它似乎是问题的答案：）