Row wise combinations

1

Is there a way to efficiently, in a vectorized manner, compute the combinations along each rows of a 2D tensor individually?

For example:

a = torch.tensor([[2, 5, 6], [7, 9, 4]])
result = torch.combinations(a, 2, dim=1)

results should look like

torch.Tensor([[[2, 5], 
                       [2, 6],
                       [5, 6]],

                      [[7, 9],
                       [7, 4],
                       [4, 9]]])

This seems like a natural thing to be able to do, but I haven’t been able to figure out a way to achieve this without looping. I’d appreciate any help and pointers. Thanks!

2

@ptrblck Do you know if there’s a way to achieve this?

3

Hi lkp411!

Yes, the idea would be to construct the combinations of a 1D vector
of indices and then index into your multidimensional tensor, swapping
the dimensions around appropriately.

Like this:

>>> import torch
>>> torch.__version__
'1.9.0'
>>> a = torch.tensor([[2, 5, 6], [7, 9, 4]])
>>> a.T[torch.combinations (torch.arange (a.shape[1]), 2)].permute (2, 0, 1)
tensor([[[2, 5],
         [2, 6],
         [5, 6]],

        [[7, 9],
         [7, 4],
         [9, 4]]])

Best.

K. Frank

4

Oh that’s clever. Thanks a lot @KFrank !!