0
votes

Is it possible to multiply only one (last) dimension in a tensor alone with other vectors?

For example, assume a tensor T=[100, 20, 400] and a matrix M =[400, 400]. Is it possible to make the operation h_{transpose}*M*h, where h is the last dimension in the tensor T? In other words, is it possible to make use of (possibly pytorch) built-in functions to get the resulting tensor of size [100, 20, 1]?

1
From the matrix M, how should the vector of shape [400, 1] chosen?kmario23

1 Answers

1
votes

I think the easiest (certainly the shortest) solution is with 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)

It basically says "for all (a, b, c, d) in bounds, multiply T[a, b, c] with M[c, d] and T[a, b, d] and accumulate it in res[a, b]".

Since einsum is implemented in terms of basic building blocks like mm, transpose etc, this could certainly be unrolled into a more "classical" solution, but right now my brain fails me at that.