Diagonal of Dot product of two large tensor

I am wondering is there a fast way to get the main diagonal of a dot product of two large tensor?

Example:
a = tensor([[0.9651, 0.5328],
[0.4087, 0.1779]])
b = tensor([[0.9567, 0.1259],
[0.7130, 0.4468]])

a * b = tensor([[0.9233, 0.0671],
[0.2914, 0.0795]])
expected: torch.diagonal(a*b) = tensor([0.9233, 0.0795])

Thanks!

Hi xdwang!

First, the example you give is not the “dot product” of two tensors,
as you refer to it in the title of the your thread. Rather, it is
element-wise multiplication of two tensors.

If indeed what you want is the diagonal of the element-wise product
of two tensors, it will be more efficient to extract the diagonals of
the two tensors first and then multiply them element-wise:

>>> import torch
>>> torch.__version__
'1.9.0'
>>> a = torch.tensor ([[0.9651, 0.5328], [0.4087, 0.1779]])
>>> b = torch.tensor ([[0.9567, 0.1259], [0.7130, 0.4468]])
>>> torch.diag (a) * torch.diag (b)
tensor([0.9233, 0.0795])

Best.

K. Frank

1 Like