How to Create a Tensor Range in PyTorch

A tensor range in PyTorch is a sequence of numbers that are stored in a one-dimensional tensor. In order to create a tensor range, you can use the torch.arange() function.

Syntax

Below is the syntax of the torch.arange() function:

torch.arange(start=0, end, step=1, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor

Parameters:

• start: the starting value for the set of points. Default: 0.
• end: the ending value for the set of points (exclusive).
• step: the gap between each pair of adjacent points. Default: 1.
• out: the output tensor (optional).
• dtype: the desired data type of returned tensor. Default: if None, uses a global default or infers from other arguments.
• layout: the desired layout of returned Tensor. Default: torch.strided.
• device: the desired device of returned tensor. Default: if None, use the current device for the default tensor type.
• requires_grad: if autograd should record operations on the returned tensor. Default: False

torch.arange() is useful for creating tensors with regular intervals or ranges of values. It can be used for indexing, slicing, looping, or generating synthetic data.

Examples

Below are some examples of using the torch.arange() function in practice.

Example 1

This example generates some simple tensor ranges:

import torch

# Create a tensor range of [0, 1, 2, 3, 4] with default start=0 and step=1
x = torch.arange(5)
print(x)
# Output: tensor([0, 1, 2, 3, 4])

# Create a tensor range of [1, 2, 3] with default step=1
y = torch.arange(1, 4)
print(y)
# Output: tensor([1, 2, 3])

# Create a tensor range of [1.0, 1.5, 2.0] with step=0.5
z = torch.arange(1.0, 2.5, 0.5)
print(z)
# Output: tensor([1.0000, 1.5000, 2.0000])

Example 2

In this example, we will specify the output tensor, the data type, and the gradient requirement of the returned tensor using keyword arguments:

import torch

a = torch.arange(2.0, dtype=torch.float64, requires_grad=True)
print(a)
# Output: tensor([0., 1.], dtype=torch.float64, requires_grad=True)

b = torch.empty(2)
torch.arange(10.0, 12.0, out=b)
print(b)
# Output: tensor([10., 11.])