张量中仅一维的矩阵向量乘法

Question

是否可以仅将张量中的一个（最后一个）维度与其他向量相乘？

例如，假设一个张量 T=[100, 20, 400] 和一个矩阵 M =[400, 400]。 是否可以进行操作h_{transpose}*M*h，其中h 是张量T 中的最后一个维度？换句话说，是否可以利用（可能是 pytorch）内置函数来获得大小为 [100, 20, 1] 的结果张量？

Answer 1

我认为最简单（当然是最短）的解决方案是使用einsum。

import torch

T = torch.randn(100, 20, 400)
M = torch.randn(400, 400)

res = torch.einsum('abc,cd,abd->ab', (T, M, T)).unsqueeze(-1)

它基本上说“对于所有(a, b, c, d)，将T[a, b, c]与M[c, d]和T[a, b, d]相乘，然后将其累加到res[a, b]中”。

由于einsum 是根据mm、transpose 等基本构建块实现的，这当然可以展开为更“经典”的解决方案，但现在我的大脑对此感到失望。