PyTorch: What is numpy.linalg.multi_dot() equivalent in PyTorch

user_12462568 0
Viewed 241

I am trying to perform matrix multiplication of multiple matrices in PyTorch and was wondering what is the equivalent of numpy.linalg.multi_dot() in PyTorch?

If there isn't one, what is the next best way (in terms of speed and memory) I can do this in PyTorch?

Code:

import numpy as np
import torch

A = np.random.rand(3, 3)
B = np.random.rand(3, 3)
C = np.random.rand(3, 3)

results = np.linalg.multi_dot(A, B, C)

A_tsr = torch.tensor(A)
B_tsr = torch.tensor(B)
C_tsr = torch.tensor(C)

# What is the PyTorch equivalent of np.linalg.multi_dot()?

Many thanks!

pytorch matrix-multiplication numpy python
2 Answers

~~Looks like one can send tensors into multi_dot~~

Looks like the numpy implementation casts everything into numpy arrays. If your tensors are on the cpu and detached this should work. Otherwise, the conversion to numpy would fail.

So in general - likely there isn't an alternative. I think your best shot is to take the multi_dot implementation, e.g. from here for numpy v1.19.0 and adjust it to handle tensors / skip the cast to numpy. Given the similar interface and the code simplicity I think that this should be pretty straightforward.

user_21319570

PyTorch has the same interface: torch.linalg.multi_dot

You can use it like:

from torch.linalg import multi_dot

multi_dot([torch.tensor([1, 2]), torch.tensor([2, 3])])
multi_dot([torch.tensor([[1, 2]]), torch.tensor([2, 3])])
multi_dot([torch.tensor([[1, 2]]), torch.tensor([[2], [3]])])

A = torch.arange(2 * 3).view(2, 3)
B = torch.arange(3 * 2).view(3, 2)
C = torch.arange(2 * 2).view(2, 2)
result = multi_dot((A, B, C))
print(result)
user_76214910
Related