Join two tensors based on a value along a dimension (similar to SQL join)

Is it possible to do an operation similar to a SQL join or SQL merge in Pytorch? For instance, I have 2 tensors A and B (both two dimensional) and I want to join them along dimension 0. In other words, I want to grab the row B[i, 1:] from tensor B and concatenate it to row A[j, :] if A[j, 0] == B[i, 0] and I want to do this for all rows of A (equivalent to SQL left join). We can assume that for each row j in A there exists only one such row i in B. Can we do this without a for loop so that we don’t slow down the gradient flow?

1 Like