Adding matrix rows to columns in numpy

Question

Say I have two 3D matrices/tensors with dimensions:

[10, 3, 1000]
[10, 4, 1000]

How do I add each combination of the third dimensions of each vector together such that to get a dimension of:

[10, 3, 4, 1000]

So each row if you will, in the second x third dimension for each of the vectors adds to the other one in every combination. Sorry if this is not clear I'm having a hard time articulating this...

Is there some kind of clever way to do this with numpy or pytorch (perfectly happy with a numpy solution, though I'm trying to use this in a pytorch context so a torch tensor manipulation would be even better) that doesn't involve me writing a bunch of nested for loops?

Nested loops example:

x = np.random.randint(50, size=(32, 16, 512))
y = np.random.randint(50, size=(32, 21, 512))
scores = np.zeros(shape=(x.shape[0], x.shape[1], y.shape[1], 512))

for b in range(x.shape[0]):
    for i in range(x.shape[1]):
        for j in range(y.shape[1]):

            scores[b, i, j, :] = y[b, j, :] + x[b, i, :]

Answer 1

Does it work for you?

import torch

x1 = torch.rand(5, 3, 6)
y1 = torch.rand(5, 4, 6)

dim1, dim2 = x1.size()[0:2], y1.size()[-2:]
x2 = x1.unsqueeze(2).expand(*dim1, *dim2)
y2 = y1.unsqueeze(1).expand(*dim1, *dim2)

result = x2 + y2

print(x1[0, 1, :])
print(y1[0, 2, :])
print(result[0, 1, 2, :])

Output :

 0.2884
 0.5253
 0.1463
 0.4632
 0.8944
 0.6218
[torch.FloatTensor of size 6]


 0.5654
 0.0536
 0.9355
 0.1405
 0.9233
 0.1738
[torch.FloatTensor of size 6]


 0.8538
 0.5789
 1.0818
 0.6037
 1.8177
 0.7955
[torch.FloatTensor of size 6]