如何实现torch.arange的tensor版本

音程2023-10-05 15:54

python 复制代码

import torch

我们都知道，torch.arange只支持数字，不支持tensor，如下：

python 复制代码

torch.arange(0,5,1)

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

但是如果使用tensor，就会报错：

python 复制代码

torch.arange(torch.tensor([0]),torch.tensor([5]),torch.tensor([1]))

可问题是，我们有如下场景怎么办:

python 复制代码

torch.arange(torch.tensor([0,2]),torch.tensor([5,7]),torch.tensor([1,1]))

也就是说，我们希望

python 复制代码

torch.arange(0,5,1)和torch.arange(2,7,1)

并行做，难道就不行吗？

上面这种并行是可以做到的，如下：

python 复制代码

x=torch.arange(0,5,1).reshape(1,-1)
a=torch.tensor([0,2])
a=a.reshape(-1,1)
x=a+x
x

tensor([[0, 1, 2, 3, 4],
$2, 3, 4, 5, 6\]\])$

细心的人可以发现，上面是具有特殊性的，

python 复制代码

torch.arange(torch.tensor([0,2]),torch.tensor([5,7]),torch.tensor([1,1]))

即

python 复制代码

torch.tensor([0,2])+5=torch.tensor([5,7])

且步长是一样的：

python 复制代码

torch.tensor([1,1])#步长都是1

为什么必须这样呢？因为这样才可以保证，输出的结果的维度是一样的。比如你换一个：

python 复制代码

torch.arange(torch.tensor([0,2]),torch.tensor([5,7]),torch.tensor([1,2]))

步长一个是1，一个是2，这样肯定不行，两个arange输出的维度不同，就肯定不可能并行。再比如：

python 复制代码

torch.arange(torch.tensor([0,2]),torch.tensor([5,9]),torch.tensor([1,1]))

步长是一样了，但是5-0=5，9-2=7（end-start），最终arange输出的维度还是会不同，无法并行。