How to do large-scale outer product efficiently

I have two matrices A and B. A and B have the same number of rows (m), and different number of columns. A.shape = [m,M], B.shape = [m,N]. For each row of A and B, I want to do an outer product (torch.outer).

Is there a way not using loop, because m is very large and I need to do this operation many times.


torch.outer only works on two vectors

This seems to work if you want to just do it for 2D tensors.

if you have a tensor A of shape [m,M] and another tensor B of shape [m,N] you can do an outer product via torch.einsum,

outer_product = torch.einsum("bi,bj->bij",A,B)