如何在PyTorch中实现三个不同维度张量乘法得到(d,d,N)张量

已知张量维度：

• L：(d, a)
• B：(a, a, N)
• R：(a, d)

目标是得到维度为(d, d, N)的结果张量，以下是两种实用实现方式：

方法一：使用torch.einsum（直观易懂） #

einsum可通过维度符号直接描述张量相乘逻辑，代码示例：

import torch

# 定义示例维度，可按需修改
d, a, N = 2, 3, 4
# 创建示例张量
L = torch.randn(d, a)
B = torch.randn(a, a, N)
R = torch.randn(a, d)

# 核心计算：按维度映射完成相乘
result = torch.einsum('da, aan, ad -> ddn', L, B, R)
print(result.shape)  # 输出: torch.Size([2, 2, 4])

'da, aan, ad -> ddn'的含义：将L的(d,a)、B的(a,a,N)、R的(a,d)按对应维度相乘并求和，最终保留(d,d,N)维度。

方法二：使用矩阵乘法torch.matmul（性能优先） #

利用PyTorch对矩阵乘法的广播支持，分步计算：

import torch

d, a, N = 2, 3, 4
L = torch.randn(d, a)
B = torch.randn(a, a, N)
R = torch.randn(a, d)

# 第一步：L与B相乘，得到(d, a, N)
LB = L @ B
# 第二步：LB与R相乘，自动识别N为批量维度，得到(d, d, N)
result = LB @ R
print(result.shape)  # 输出: torch.Size([2, 2, 4])

这里@是torch.matmul的简写，PyTorch会自动处理批量维度的矩阵乘法优化。

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