1

Supposing I have two square matrices A, B of the same size

A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[1, 1], [1, 1]])

And I want a resulting tensor that consists of the row-wise dot product, say

tensor([3, 7])  # i.e. (1*1 + 2*1, 3*1 + 4*1)

What is an efficient means of achieving this in PyTorch?

Ivan
  • 34,531
  • 8
  • 55
  • 100
IntegrateThis
  • 853
  • 2
  • 16
  • 39

2 Answers2

1

As you said you can use torch.bmm but you first need to broadcast your inputs:

>>> torch.bmm(A[..., None, :], B[..., None])
tensor([[[3]],

        [[7]]])

Alternatively you can use torch.einsum:

>>> torch.einsum('ij,ij->i', A, B)
tensor([3, 7])
Ivan
  • 34,531
  • 8
  • 55
  • 100
  • Thanks for the answer. Do you have any good references for understanding the einsum notation, I feel I'm a bit confused even reading the Numpy documentation about it. Also, I'm assuming the None values in the slice notation do the same thing as unsqueeze at a given axis location? – IntegrateThis Oct 25 '21 at 06:50
  • 1
    @IntegrateThis You can look into [this post](https://stackoverflow.com/questions/55894693/understanding-pytorch-einsum): it gives some intuition on the einsum interface as well as examples and a possible implementation. The post is about `np.einsum` but the same can be said for `torch.einsum` as they both share roughly the same features. Indeed indexing with `None` allows to unsqueeze a additional dimension, just like you did with the reshape. – Ivan Oct 25 '21 at 07:31
0
import torch
import numpy as np

def row_wise_product(A, B):
    num_rows, num_cols = A.shape[0], A.shape[1]
    prod = torch.bmm(A.view(num_rows, 1, num_cols), B.view(num_rows, num_cols, 1))
    return prod

A = torch.tensor(np.array([[1, 2], [3, 4]]))
B = torch.tensor(np.array([[1, 1], [1, 1]]))
C = row_wise_product(A, B)
IntegrateThis
  • 853
  • 2
  • 16
  • 39