Outer product for batched and unbatched data

Hello everyone.

In my code I’ll need to perform the outer product for both unbatched and batched data (tensors with dimension T1(Ndim) and T2(B, Ndim) respectively).

For now I was using:

prod = T1[:, None] * T1 # Unbatched
prod = T2[:, :, None] * T2[:, None, :] # Batched

But, is there a way to have an unified expression for both computations?


Well I found the solution. This works for both cases:

prod = T[..., None, :] * T[..., None]

where T can be either T1 or T2.