如何为einsum运算添加batch维度？TensorFlow实现验证

你的Einstein求和表达式修改是正确的，完全符合需求。

逻辑说明 #

原无batch维度时，表达式'ijh, jhu -> ihu'的核心是对两个张量的j维度做求和收缩，最终保留i（目标节点）、h（注意力头）、u（特征维度）。

添加batch维度后：

• 输入张量的batch维度b是两个张量的首个维度，仅作为独立批次标识，不参与任何维度收缩
• 表达式'bijh, bjhu -> bihu'中，b在输入和输出中都保留，其余维度的收缩逻辑和原表达式完全一致：对j维度求和，最终得到(batch_size, n_nodes, n_heads, n_units)的输出，完全匹配需求。

修改后的完整代码示例 #

import tensorflow as tf

tf.random.set_seed(0)

# 添加batch维度，示例batch_size=3
tensor = tf.random.uniform(shape=(3, 2, 2, 2))  # Shape: (batch_size, n_nodes, n_nodes, n_heads)
tensor1 = tf.random.uniform(shape=(3, 2, 2, 2)) # Shape: (batch_size, n_nodes, n_heads, n_units)

print("输入tensor形状:", tensor.shape)
print("输入tensor1形状:", tensor1.shape)
print("-" * 50)

einsum_tensor = tf.einsum('bijh, bjhu -> bihu', tensor, tensor1)

print("输出einsum_tensor形状:", einsum_tensor.shape)
print(einsum_tensor)

运行后会输出形状为(3, 2, 2, 2)的结果，符合预期。

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