end, start and step argument of arange() work with a 0D tensor against error messages

[hyperkai]

🐛 Describe the bug

end argument of arange() cannot have a tensor according to the error message as shown below:

import torch
                            # ↓↓↓↓↓↓↓↓↓
torch.arange(end=torch.tensor([0, 1, 2])) # Error

TypeError: arange(): argument 'end' must be Number, not Tensor

And, start argument of arange() cannot have a tensor according to the error message as shown below:

import torch
                              # ↓↓↓↓↓↓↓↓↓
torch.arange(start=torch.tensor([0, 1, 2]), end=5) # Error

TypeError: arange() received an invalid combination of arguments - got (end=int, start=Tensor, ), but expected one of:
 * (Number end, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (Number start, Number end, *, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (Number start, Number end, Number step, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)

And, step argument of arange() cannot have a tensor according to the error message as shown below:

import torch
                                             # ↓↓↓↓↓↓↓↓↓
torch.arange(start=0, end=5, step=torch.tensor([0, 1, 2])) # Error

TypeError: arange() received an invalid combination of arguments - got (end=int, step=Tensor, start=int, ), but expected one of:
 * (Number end, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (Number start, Number end, *, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (Number start, Number end, Number step, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)

But, end, start and step argument of arange() work with the 0D tensor of `an integer against the error messages as shown below:

import torch
                            # ↓
torch.arange(end=torch.tensor(2)) # tensor([0, 1])
                              # ↓
torch.arange(start=torch.tensor(2), end=5) # tensor([2, 3, 4])
                                             # ↓
torch.arange(start=0, end=5, step=torch.tensor(2)) # tensor([0, 2, 4])

In addition, , end, start and step argument of arange() work with the 0D tensor of a floating-point number, complex number or boolean value against the error messages as shown below:

import torch
                            # ↓↓
torch.arange(end=torch.tensor(2.)) # tensor([0., 1.])
                            # ↓↓↓↓↓↓
torch.arange(end=torch.tensor(2.+0.j)) # tensor([0., 1.])
                            # ↓↓↓↓
torch.arange(end=torch.tensor(True)) # tensor([0])

                              # ↓↓
torch.arange(start=torch.tensor(2.), end=5.) # tensor([2., 3., 4.])
                              # ↓↓↓↓↓↓
torch.arange(start=torch.tensor(2.+0.j), end=5.+0.j) # tensor([2., 3., 4.])
                              # ↓↓↓↓↓
torch.arange(start=torch.tensor(False), end=True) # tensor([0])

                                               # ↓↓
torch.arange(start=0., end=5., step=torch.tensor(2.)) # tensor([0., 2., 4.])
                                                       # ↓↓↓↓↓↓
torch.arange(start=0.+0.j, end=5.+0.j, step=torch.tensor(2.+0.j)) # tensor([0., 2., 4.])
                                                    # ↓↓↓↓
torch.arange(start=False, end=True, step=torch.tensor(True)) # tensor([0])

Versions

import torch

torch.__version__ # 2.3.0+cu121

cc @gchanan @mruberry

[janeyx99]

This likely has to do with the fact that the dispatcher will reinterpret those Scalar Tensors as Scalars, which is a good feature for other operations in PyTorch. (cc @bdhirsh to confirm.)

The current behavior seems fine to me--is there an action item to pursue here?

[malfet]

Closing, as indeed it works as expected: ability to treat scalar tensors as numbers fits nicely with Python's duck-typing paradigm