PyTorch 创建范围组合

Question

有没有一种方法可以轻松嵌套arange以在 PyTorch 中创建两个范围的所有组合？ 例如：

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]]

即，与此 python 代码具有相同功能的东西：

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

我们可以更改range(x,y)的位置。

Answer 1

考虑以下：

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])

结果：

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

Answer 2

幸运的是，有一个内置函数可以满足您的需要，它就是torch.cartesian_prod()函数。

这是一个例子：

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