如何实现torch.arange的tensor版本

文章目录

• • • 背景
    • 实现方案
    • 不可行的情况

背景

import torch

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

torch.arange(0,5,1)

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

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

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

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

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

也就是说，我们希望

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

并行做，难道就不行吗？

实现方案

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

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

不可行的情况

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

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

即

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

且步长是一样的：

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

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

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

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

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

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