multiply many matrices and many vectors pytorch

user_9704615 0
Viewed 44

I am trying to multiply the following:

A batch of matrices N x M x D
A batch of vectors N x D x 1
To get a result: N x M x 1

as if I were doing N dot products on M x D D x 1.

I cant seem to find the correct function in PyTorch.

torch.bmm as far as I can tell only works for a batch of vectors and a single matrix. If I have to use torch.einsum then so be it but id rather not!

pytorch linear-algebra python-3.x matrix
1 Answers

It's pretty straightforward and intuitive with einsum:

torch.einsum('ijk, ikl->ijl', mats, vecs)

But your operation is just:

mats @ vecs
user_42384080
Related