Pytorch broadcasting product of two tensors

Question

I want to multiply two tensors, here is what I have got:

A tensor of shape (20, 96, 110)
B tensor of shape (20, 16, 110)

The first index is for batch size. What I want to do is essentially take each tensor from B - (20, 1, 110), for example, and with that, I want to multiply each A tensor (20, n, 110). So the product will be at the end: tensor AB which shape is (20, 96 * 16, 110).

So I want to multiply each tensor from A by broadcasting with B. Is there a method in PyTorch that does it?

Im not sure, but from what i can gather because you want to multiply the first element of tensor a to the first element of tensor b,The product will be always the same for all the 20 batches, So why dont you just plainly initialize a tensor of torch.tensor(20, 96*16, 110) here 20 being the batch size — Ryan

benjaminplanche benjaminplanche · Accepted Answer · 2018-06-22T10:10:07

Using torch.einsum followed by torch.reshape:

AB = torch.einsum("ijk,ilk->ijlk", (A, B)).reshape(A.shape[0], -1, A.shape[2])

Example:

import numpy as np
import torch

# A of shape (2, 3, 2):
A = torch.from_numpy(np.array([[[1, 1], [2, 2], [3, 3]], 
                               [[4, 4], [5, 5], [6, 6]]]))
# B of shape (2, 2, 2):
B = torch.from_numpy(np.array([[[1, 1], [10, 10]], 
                               [[2, 2], [20, 20]]]))

# AB of shape (2, 3*2, 2):
AB = torch.einsum("ijk,ilk->ijlk", (A, B)).reshape(A.shape[0], -1, A.shape[2])
# tensor([[[ 1, 1], [ 10, 10], [  2,  2], [ 20,   20], [ 3,   3], [ 30,  30]],
#         [[ 8, 8], [ 80, 80], [ 10, 10], [ 100, 100], [ 12, 12], [ 120, 120]]])

Pytorch broadcasting product of two tensors

1 Answers