I’m looking a way to do the following thing: Let’s assume we have a tensor A of dimension [N,F] and a tensor B of dimension [N,F], I would like to obtain a tensor C of dimension [N,N,2*F].

It wasn’t a typo, I indeed would like a tensor NxNx2F. Otherwise, cat would have been perfect However, I’m not aware of something like this. I thought about something like torch.repeat but didn’t find exactly what I’m looking for

I need to implement a function given as input x: 1xF, y: 1xF and outputs a scalar. I’m looking for to implement it in a GPU efficient manner, i.e. create all pair x_i, y_j and compute a scalar with one operation. This is why I though about this (NxF)x(NxF) --> NxNx2F

Yes exactly. Why do you think it would be more appropriate ?
My idea afterwards was to use a feed forward net to compute the attention coefficient, using 2F as input.

Well, if x contains F elements and y contains F elements then there are FxF pairings, not 2xF pairings.

But, I had forgotten that there are two elements in each pairing. So the size should be NxFxFx2.

I can’t see why you need the pairings but this should do it.

A = A.unsqueeze(2).repeat([1, 1, F]) # shape NxFxF with every value repeated along the last dim
B = B.unsqueeze(1).repeat([1, F, 1]) # shape NxFxF with every value repeated along the middle dim
C = torch.stack((A, B), dim=3)

C has 4 dimensions of size N, F, F, 2
C[n, fb, fa] contains the pairing A[fa], B[fb] from batch n.

Thank you for your answer, I’ll have a try later if it might fit. I see your point. The 2xF pairing it’s because I would like to concatenate the last dimension in order to have 2F and simply use a weight matrix 2Fx1 instead of FxFx1.

2xF is enough to store all elements of both A and B, but not enough to explicitly store all the pairings. That said, if you are just going to add a linear layer in order to produce one scalar of output, then 2xF should be fine, and this will do the job

C = torch.stack([A,B], dim=2)
C_2F = C.view(N, -1)