# PyTorch Add Dimension: Expanding a Tensor with a Dummy Axis

Ben Cook • Posted 2017-03-09 • Last updated 2021-10-21 • Code

Adding a dimension to a tensor can be important when you’re building machine learning models. Although the actual PyTorch function is called `unsqueeze()`, you can think of this as the PyTorch “add dimension” operation. Let’s look at two ways to do it.

## Using `None` indexing

The easiest way to expand tensors with dummy dimensions is by inserting `None` into the axis you want to add. For example, say you have a feature vector with 16 elements. To add a dummy batch dimension, you should index the 0th axis with `None`:

``````import torch

x = torch.randn(16)
x = x[None, :]
x.shape

# Expected result
# torch.Size([1, 16])``````

The slicing syntax works by specifying new dimensions with `None` and existing dimensions with a colon. That means you can prepend more than one dimension if you want:

``````x = torch.randn(16)
x = x[None, None, :]
x.shape

# Expected result
# torch.Size([1, 1, 16])``````

Additionally, when you prepend indices, you don’t actually have to specify the colons. All the remaining dimensions are implied. That means you can repeat the same example with no colon at the end:

``````x = torch.randn(16)
x = x[None, None]
x.shape

# Expected result
# torch.Size([1, 1, 16])``````

This is a very simple trick for prepending axes to the front of tensors.

You can also add dimensions to the end of tensors, which can be useful for broadcasting operations like pairwise distance. All you have to do is rearrange the colons and the `None`(s).

Let’s convert a 2D tensor with shape `(3, 4)` to a 3D tensor with shape `(3, 4, 1)`:

``````x = torch.randn(3, 4)
x = x[:, :, None]
x.shape

# Expected result
# torch.Size([3, 4, 1])``````

Notice that we have to specify the first two dimensions with colons and then add a `None` to the end. When you append dimensions to the end, the colons are required.

And naturally, this trick works regardless of where you want to insert the dimension. For example, if you wanted to add a dimension to the middle of the previous tensor, you could write `x[:, None, :]`.

Overall, I prefer using `None` to add dimensions because the same syntax works in NumPy:

``````import numpy as np

x = np.random.randn(16)
x = x[None]
x.shape

# Expected result
# (1, 16)``````

## PyTorch Unsqueeze

But there is another approach that you should know about. PyTorch provides a function called `unsqueeze()` that does the same thing.

``````x = torch.randn(16)
x = torch.unsqueeze(x, dim=0)
x.shape

# Expected result
# torch.Size([1, 16])``````

The `dim` argument is how you specify where the new axis should go. To put a new dimension on the end, pass `dim=-1`:

``````x = torch.randn(3, 4)
x = torch.unsqueeze(x, dim=-1)
x.shape

# Expected result
# torch.Size([3, 4, 1])``````

Not bad. But you have to be careful if you use both NumPy and PyTorch because there is no NumPy `unsqueeze()` function:

``````x = np.random.randn(16)
try:
x = np.unsqueeze(x, axis=0)
except AttributeError as e:
print(e)

# Expected error message
# module 'numpy' has no attribute 'unsqueeze'``````

Indeed, the equivalent NumPy function is called `expand_dims()`:

``````x = np.random.randn(16)
x = np.expand_dims(x, axis=0)
x.shape

# Expected error
# (1, 16)``````

Better to use `None` indexing for the PyTorch “add dimension” op to avoid confusion.