Calculating joint probabilities from two tensors of different sizes in pytorch

Question

I am trying to calculate joint probabilities from two tensors.. It's a little bit confusing for me. Suppose we have :

a = torch.Tensor((10, 2))
b = torch.Tensor((10, 5))

c would be then of size (10, 5*2) = (10,10) I want to calculate c such as , e.g for the first row:

c[0,0] = a[0,0] * b[0,0]
c[0,1] = a[0,0] * b[0,1]
c[0,2] = a[0,0] * b[0,2]
...
c[0,5] = a[0,0] * b[0,5]
c[0,6] = a[0,1] * b[0,0]
c[0,7] = a[0,1] * b[0,1]
....
c[0,10] = a[0,1] * b[0,5]



   

Answer 1

I guess you assume a, b to be independent. Hence joint probability is p(a)p(b). With that, you can use the code below. Explanations are in the code comments.

import torch

# variables to work with
a = torch.rand((10,2))
b = torch.rand((10,5))


# joint is the ground truth and we compute it using for loop 
# as described in the question. 
joint = torch.rand((10,10))

for i in range(2):
    for j in range(5):
        joint[:, i*5+j] = a[:, i]*b[:, j]

#
# using outerproduct and flattening
# 
# if we had two vector of probabilities, 
#    the outer product will give us elementwise multiplications. 
#    We can then flatten it to get p(a,b). 
#    We use the same idea here

c = a.unsqueeze(2) @ b.unsqueeze(1)
c = c.reshape(10,-1)

# see if we got the correct result
print(torch.allclose(c,joint))