2D times 2D equals a 3d pytorch tensor

Question

Given two 2-D pytorch tensors:

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

Is there an efficient way to calculate a tensor of shape (6, 2, 2) where each entry is a column of A times each row of B?

For example, with A and B above, the 3D tensor should have the following matrices:

[[[0, 0],
  [0, 0]],

 [[1, 1],
  [3, 3]],

 [[2, 2],
  [6, 6]],

 [[0, 0],
  [0, 0]],

 [[2, 2],
  [4, 4]],

 [[4, 4],
  [8, 8]]]

I know how to do it via for-loop but I am wondering if could have an efficient way to save it.

Do you mean matrix multiplication or broadcasted element-wise multiplication? Both give the same results when the rows and columns of A and B are length 2. — jodag

jodag jodag · Accepted Answer · 2021-01-19T04:17:24

Pytorch tensors implement numpy style broadcast semantics which will work for this problem.

It's not clear from the question if you want to perform matrix multiplication or element-wise multiplication. In the length 2 case that you showed the two are equivalent, but this is certainly not true for higher dimensionality! Thankfully the code is almost the same so I'll just give both options.

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

# matrix multiplication
C_mm = (A.T[:, None, :, None] @ B[None, :, None, :]).flatten(0, 1)

# element-wise multiplication
C_ew = (A.T[:, None, :, None] * B[None, :, None, :]).flatten(0, 1)

Code description. A.T transposes A and the indexing with None inserts unitary dimensions so A.T[:, None, :, None] will be shape (2, 1, 2, 1) and B[None, :, None, :] is shape (1, 3, 1, 2). Since @ (matrix multiplication) operates on the last two dimensions of tensors, and broadcasts the other dimensions, then the result is matrix multiplication for each column of A times each row of B. In the element-wise case the broadcasting is performed on every dimension. The result is a (2, 3, 2, 2) tensor. To turn it into a (6, 2, 2) tensor we just flatten the first two dimensions using Tensor.flatten.

2D times 2D equals a 3d pytorch tensor

1 Answers