CODEWITHSUNDEEP
pytorch
IntegrateThis
IntegrateThis

Reputation: 962

Efficient method to compute the row-wise dot product of two square matrices of the same size in PyTorch

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?

Upvotes: 1

Views: 989

Answers (2)

Ivan
Ivan

Reputation: 40618

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])

Upvotes: 1

IntegrateThis
IntegrateThis

Reputation: 962

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)

Upvotes: 0

Related Questions