Indexing a tensor or ndarray with `None`

Last updated Apr 20, 2022 Edit Source

• All
• Numpy
• Py torch

Tags: #Numpy #PyTorch

Similar to NumPy you can insert a singleton dimension (“unsqueeze” a dimension) by indexing this dimension with None. In turn n[:, None] will have the effect of inserting a new dimension on dim=1. This is equivalent to n.unsqueeze(dim=1):

1
2
3
4
5
6
7

>>> n = torch.rand(3, 100, 100)

>>> n[:, None].shape
(3, 1, 100, 100)

>>> n.unsqueeze(1).shape
(3, 1, 100, 100)

In the example above : is was used as a placeholder to designate the first dimension dim=0. If you want to insert a dimension on dim=2, you can add a second : as n[:, :, None].

You can also place None with respect to the last dimension instead. To do so you can use the ellipsis syntax ...:

• n[..., None] will insert a dimension last, i.e. n.unsqueeze(dim=-1).

• n[..., None, :] on the before last dimension, i.e. n.unsqueeze(dim=-2).

None is a version with advanced indexing , which might be a bit slower because it has more checking to do to find out exactly what you want to do.

Sources:

• syntax - Indexing a tensor with None in PyTorch - Stack Overflow
• What is the difference between [None, …] and unsqueeze? - PyTorch Forums