PyTorch create combination of ranges-CodePudding

Is there a way to nest arange's easily to create all the combinations of two ranges in PyTorch? For example:

x = torch.arange(2, 4)
y = torch.arange(0, 3)

something(x, y)
# should be [[2,0], [2,1], [2,3], [3,0], [3,1], [3,2]]

I.e., something with the same functionality as this python code:

l = []
for x in range(2, 4):
   for y in range(0, 3):
       l.append([x, y])

where the range(x,y) we can change.

CodePudding user response：

Fortunately there's a built-in function that does what you need and it is the torch.cartesian_prod() function.

Here's an example:

x = torch.arange(2, 4)
y = torch.arange(0, 3)
l = torch.cartesian_prod(x,y)

CodePudding user response：

Consider the following:

x = torch.arange(2, 4)
y = torch.arange(0, 3)

m,n = len(x),len(y)
res = torch.stack([
    x[:,None].broadcast_to(m,n),
    y.broadcast_to(m,n)]
    ).permute(1,2,0).reshape([m*n,2])

Result:

tensor([[2, 0],
        [2, 1],
        [2, 2],
        [3, 0],
        [3, 1],
        [3, 2]])