两个3D张量执行tensordot运算的维度设置：如何得到B×D形状输出

可以完全通过向量化运算实现目标输出，无需遍历B维度，以下是不同实现方案：

运算逻辑说明 #

你的需求本质是对每个batch下的样本，用形状为1×D的张量和形状为L×D的张量做对应维度的乘加运算，最终保留B×D的输出维度。你之前调用tensordot出现多余B维度的核心原因是：原生tensordot会将两个输入未参与收缩的所有维度按顺序拼接，默认不会自动对齐两个输入的公共B维度，因此会输出B×D×B×D的冗余结果。

方案1：广播求和（最简洁高效） #

直接利用张量广播特性实现，代码可读性和运算效率都优于硬调tensordot参数：

# numpy 示例
import numpy as np
B, L, D = 32, 10, 64
# 定义输入张量
X = np.random.randn(B, L, D) # 形状 B×L×D
Y = np.random.randn(B, 1, D) # 形状 B×1×D
# 广播相乘后对L维度求和
result = (X * Y).sum(axis=1)
print(result.shape) # 输出 (32, 64) 即 B×D

PyTorch 实现逻辑完全一致：

import torch
B, L, D = 32, 10, 64
X = torch.randn(B, L, D)
Y = torch.randn(B, 1, D)
result = (X * Y).sum(dim=1)
print(result.shape) # 输出 torch.Size([32, 64])

方案2：einsum 实现（维度逻辑最清晰） #

如果需要明确控制每个维度的运算规则，einsum是最优选择，不需要调整张量维度顺序：

# numpy 版本
result = np.einsum("bld,b1d->bd", X, Y)
# PyTorch 版本
result = torch.einsum("bld,b1d->bd", X, Y)

如果你需要得到之前循环输出的B×L形状结果（即在D维度做乘加），只需要修改输出标识即可：

result = np.einsum("bld,b1d->bl", X, Y)

方案3：tensordot 实现 #

tensordot更适合不需要保留公共批量维度的收缩场景，如果你一定要用tensordot实现，可以通过维度重组的方式完成：

# 先将Y的维度调整为 B×D×1，再对L和1维度做收缩
Y_trans = Y.transpose(0,2,1)
# 收缩X的1轴（L）和Y_trans的2轴（1维度），再对结果做维度压缩
result = np.tensordot(X, Y_trans, axes=([1], [2])).diagonal(axis1=0, axis2=2).transpose(1,0)
print(result.shape) # 输出 (32,64)

该方案可读性较差，更推荐使用前两种方案。

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