如何在PyTorch中生成具有不同向量的新Tensor？

Question

I want to generate new a○b vector with a and b (○ means element wise multiply). My code is below, but the performance looks bad because of for . Are there any efficient way?

a = torch.rand(batch_size, a_len, hid_dim)
b = torch.rand(batch_size, b_len, hid_dim)
# a_elmwise_mul_b = torch.zeros(batch_size, a_len, b_len, hid_dim)
for sample in range(batch_size):
    for ai in range(a_len):
        for bi in range(b_len):
            a_elmwise_mul_b[sample, ai, bi] = torch.mul(a[sample, ai], b[sample, bi])

Update

I updated my code refer to Ahmad! Thank you.

N = 16
hid_dim = 50
a_seq_len = 10
b_seq_len = 20
a = torch.randn(N, a_seq_len, hid_dim)
b = torch.randn(N, b_seq_len, hid_dim)
shape = (N, a_seq_len, b_seq_len, hid_dim)

a_dash = a.unsqueeze(2) # (N, a_len, 1,     hid_dim)
b_dash = b.unsqueeze(1) # (N, 1,     b_len, hid_dim)
a_dash = a_dash.expand(shape)
b_dash = b_dash.expand(shape)
print(a_dash.size(), b_dash.size())
mul = a_dash * b_dash
print(mul.size())
----------
torch.Size([16, 10, 20, 50]) torch.Size([16, 10, 20, 50])
torch.Size([16, 10, 20, 50])

Answer 1

From your problem definition, it looks like you want to multiply two tensors, say A and B of shape AxE and BxE and want to get a tensor of shape AxBxE . It means you want to multiply, each row of tensor A with the whole tensor B . If it is correct, then we don't call it element-wise multiplication.

You can accomplish your goal as follows.

import torch

# batch_size = 16, a_len = 10, b_len = 20, hid_dim = 50
a = torch.rand(16, 10, 50)
b = torch.rand(16, 20, 50)

c = a.unsqueeze(2).expand(*a.size()[:-1], b.size(1), a.size()[-1])
d = b.unsqueeze(1).expand(b.size()[0], a.size(1), *b.size()[1:])
print(c.size(), d.size())
print(c.size(), d.size())

mul = c * d       # shape of c, d: 16 x 10 x 20 x 50
print(mul.size()) # 16 x 10 x 20 x 50

Here, mul tensor is your desired result. Just to clarify, the above two lines realted to c and d computation, are equivalent to:

c = a.unsqueeze(2).expand(a.size(0), a.size(1), b.size(1), a.size(2))
d = b.unsqueeze(1).expand(b.size(0), a.size(1), b.size(1), b.size(2))