note.nkmk.me

NumPy: Add new dimensions to an array (np.newaxis, np.expand_dims)

Modified: | Tags: Python, NumPy

In NumPy, you can add new dimensions to an array (ndarray) using np.newaxis, np.expand_dims(), and np.reshape().

Contents

np.expand_dims() is similar to torch.unsqueeze() in PyTorch. Note that np.unsqueeze() is not provided.

The np.reshape() function or the reshape() method of ndarray can not only add dimensions but also transform to any desired shape. For details, refer to the following article.

To remove dimensions of size 1, use np.squeeze().

The NumPy version used in this article is as follows. Note that functionality may vary between versions.

import numpy as np

print(np.__version__)
# 1.26.1
source: numpy_newaxis.py

How to use np.newaxis

np.newaxis is None

np.newaxis is an alias for None.

print(np.newaxis is None)
# True
source: numpy_newaxis.py

It is just named differently for clarity, so np.newaxis in the following sample code can be replaced with None and it will work the same.

Add new dimensions with np.newaxis

Using np.newaxis within [] adds a new dimension of size 1 at that position.

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(a.shape)
# (2, 3)

print(a[:, :, np.newaxis])
# [[[0]
#   [1]
#   [2]]
# 
#  [[3]
#   [4]
#   [5]]]

print(a[:, :, np.newaxis].shape)
# (2, 3, 1)
source: numpy_newaxis.py
print(a[:, np.newaxis, :])
# [[[0 1 2]]
# 
#  [[3 4 5]]]

print(a[:, np.newaxis, :].shape)
# (2, 1, 3)
source: numpy_newaxis.py
print(a[np.newaxis, :, :])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis, :, :].shape)
# (1, 2, 3)
source: numpy_newaxis.py

The trailing : within [] can be omitted. To add a dimension at the beginning, just [np.newaxis] is sufficient.

print(a[:, np.newaxis])
# [[[0 1 2]]
# 
#  [[3 4 5]]]

print(a[:, np.newaxis].shape)
# (2, 1, 3)
source: numpy_newaxis.py
print(a[np.newaxis])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis].shape)
# (1, 2, 3)
source: numpy_newaxis.py

Consecutive : can be replaced with ..., which is particularly useful for adding dimensions to higher-dimensional ndarray.

print(a[..., np.newaxis])
# [[[0]
#   [1]
#   [2]]
# 
#  [[3]
#   [4]
#   [5]]]

print(a[..., np.newaxis].shape)
# (2, 3, 1)
source: numpy_newaxis.py

You can use multiple np.newaxis at once. Multiple dimensions are added.

print(a[np.newaxis, :, np.newaxis, :, np.newaxis])
# [[[[[0]
#     [1]
#     [2]]]
# 
# 
#   [[[3]
#     [4]
#     [5]]]]]

print(a[np.newaxis, :, np.newaxis, :, np.newaxis].shape)
# (1, 2, 1, 3, 1)
source: numpy_newaxis.py

The view of the original object is returned when adding dimensions with np.newaxis. Since the original object and the view object share memory, changing a value in one could affect the other.

a_newaxis = a[:, :, np.newaxis]

print(np.shares_memory(a, a_newaxis))
# True
source: numpy_newaxis.py

Control broadcasting with np.newaxis

In binary operations (such as arithmetic operations) between NumPy arrays, they are automatically reshaped into the same shape by broadcasting.

a = np.zeros(27, dtype=np.int64).reshape(3, 3, 3)
print(a)
# [[[0 0 0]
#   [0 0 0]
#   [0 0 0]]
# 
#  [[0 0 0]
#   [0 0 0]
#   [0 0 0]]
# 
#  [[0 0 0]
#   [0 0 0]
#   [0 0 0]]]

print(a.shape)
# (3, 3, 3)

b = np.arange(9).reshape(3, 3)
print(b)
# [[0 1 2]
#  [3 4 5]
#  [6 7 8]]

print(b.shape)
# (3, 3)

print(a + b)
# [[[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]]
source: numpy_newaxis.py

In broadcasting, a new dimension is added to the beginning of the array with fewer dimensions to align the number of dimensions between arrays.

Adding a new dimension at the beginning yields the same result as automatic broadcasting. 

print(b[np.newaxis, :, :].shape)
# (1, 3, 3)

print(a + b[np.newaxis, :, :])
# [[[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]
# 
#  [[0 1 2]
#   [3 4 5]
#   [6 7 8]]]
source: numpy_newaxis.py

Changing the position of the added dimension can lead to different outcomes.

print(b[:, np.newaxis, :].shape)
# (3, 1, 3)

print(a + b[:, np.newaxis, :])
# [[[0 1 2]
#   [0 1 2]
#   [0 1 2]]
# 
#  [[3 4 5]
#   [3 4 5]
#   [3 4 5]]
# 
#  [[6 7 8]
#   [6 7 8]
#   [6 7 8]]]
source: numpy_newaxis.py
print(b[:, :, np.newaxis].shape)
# (3, 3, 1)

print(a + b[:, :, np.newaxis])
# [[[0 0 0]
#   [1 1 1]
#   [2 2 2]]
# 
#  [[3 3 3]
#   [4 4 4]
#   [5 5 5]]
# 
#  [[6 6 6]
#   [7 7 7]
#   [8 8 8]]]
source: numpy_newaxis.py

For example, directly adding or subtracting a color image array (shape: (height, width, color)) with a monochrome image array (shape: (height, width)) is not possible due to incompatible shapes for broadcasting. However, adding a new dimension at the end of the monochrome image enables successful broadcasting. For more on this, refer to the following article.

Add new dimensions with np.expand_dims()

You can also add a new dimension to a NumPy array (ndarray) with np.expand_dims().

Specify the original ndarray as the first argument a and the position to add the dimension as the second argument axis.

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(np.expand_dims(a, 0))
# [[[0 1 2]
#   [3 4 5]]]

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)
source: numpy_expand_dims.py

You can insert a new dimension at any position.

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)

print(np.expand_dims(a, 1).shape)
# (2, 1, 3)

print(np.expand_dims(a, 2).shape)
# (2, 3, 1)
source: numpy_expand_dims.py

Negative values for the second argument (axis) enable reverse-order specification, with -1 denoting the last dimension.

print(np.expand_dims(a, -1).shape)
# (2, 3, 1)

print(np.expand_dims(a, -2).shape)
# (2, 1, 3)

print(np.expand_dims(a, -3).shape)
# (1, 2, 3)
source: numpy_expand_dims.py

Up to NumPy version 1.17, specifying out-of-range values for the second argument (axis) did not cause an error, allowing dimensions to be added at the end or beginning. However, starting from NumPy 1.18, such specifications result in an error.

# print(np.expand_dims(a, 3).shape)
# AxisError: axis 3 is out of bounds for array of dimension 3

# print(np.expand_dims(a, -4).shape)
# AxisError: axis -4 is out of bounds for array of dimension 3
source: numpy_expand_dims.py

From NumPy 1.18 onwards, it also became possible to specify multiple positions with a tuple as the second argument (axis) to add multiple dimensions at once.

print(np.expand_dims(a, (0, 1, -1)).shape)
# (1, 1, 2, 3, 1)
source: numpy_expand_dims.py

Like np.newaxis, np.expand_dims() also returns a view.

a_expand_dims = np.expand_dims(a, 0)

print(np.shares_memory(a, a_expand_dims))
# True
source: numpy_expand_dims.py

Examples are omitted here, but as mentioned above, it is also possible to control broadcasting by adding new dimensions with np.expand_dims().

Add new dimensions with np.reshape()

You can reshape an ndarray using the np.reshape() function or the reshape() method. See the following article for details.

By specifying a new shape with additional dimensions in reshape(), you can achieve the same result as when adding dimensions with np.newaxis or np.expand_dims().

a = np.arange(6).reshape(2, 3)
print(a)
# [[0 1 2]
#  [3 4 5]]

print(a.shape)
# (2, 3)

print(a[np.newaxis])
# [[[0 1 2]
#   [3 4 5]]]

print(a[np.newaxis].shape)
# (1, 2, 3)

print(np.expand_dims(a, 0))
# [[[0 1 2]
#   [3 4 5]]]

print(np.expand_dims(a, 0).shape)
# (1, 2, 3)

print(a.reshape(1, 2, 3))
# [[[0 1 2]
#   [3 4 5]]]

print(a.reshape(1, 2, 3).shape)
# (1, 2, 3)
source: numpy_newaxis.py

Using np.newaxis or np.expand_dims() has the advantage that there is no need to explicitly specify the sizes of the other dimensions.

Related Categories

Related Articles