np.expand_dims() and np.squeeze() in NumPy - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of np.expand_dims() on a 1D array

What is the output shape of the array after applying np.expand_dims() on a 1D array along axis 0?

NumPy

import numpy as np
arr = np.array([1, 2, 3])
expanded = np.expand_dims(arr, axis=0)
print(expanded.shape)

A (3,)

B (3, 1)

C (1, 3)

D (1,)

💡 Hint

Think about adding a new dimension at the start of the shape.

The original array has shape (3,). Using np.expand_dims with axis=0 adds a new dimension at the front, resulting in shape (1, 3).

❓ Predict Output

intermediate

Effect of np.squeeze() on a 3D array

Given a 3D array with shape (1, 4, 1), what is the shape after applying np.squeeze()?

NumPy

import numpy as np
arr = np.zeros((1, 4, 1))
squeezed = np.squeeze(arr)
print(squeezed.shape)

A (1, 4, 1)

B (4,)

C (4, 1)

D (1, 4)

💡 Hint

np.squeeze removes all dimensions of size 1.

The original shape is (1, 4, 1). np.squeeze removes all dimensions with size 1, leaving only the dimension with size 4, so the shape becomes (4,).

🔧 Debug

advanced

Identify the error in np.expand_dims usage

Which option will cause an error when trying to expand dimensions of a 2D array with shape (3, 4) using np.expand_dims()?

NumPy

import numpy as np
arr = np.ones((3, 4))
expanded = np.expand_dims(arr, axis=3)
print(expanded.shape)

A axis=4

B axis=3

C axis=-1

D axis=2

💡 Hint

Check if the axis value is valid for the array's dimensions.

The array has 2 dimensions. Valid axis values for np.expand_dims are from -3 to 2 (inclusive). axis=4 is out of range and causes an IndexError.

❓ data_output

advanced

Result of squeezing with specific axis

What is the shape of the array after applying np.squeeze() with axis=0 on an array of shape (1, 5, 1)?

NumPy

import numpy as np
arr = np.zeros((1, 5, 1))
squeezed = np.squeeze(arr, axis=0)
print(squeezed.shape)

A (5, 1)

B (1, 5)

C (5,)

D (1, 5, 1)

💡 Hint

Only the specified axis with size 1 is removed.

Using axis=0 removes the first dimension because it has size 1, resulting in shape (5, 1). The last dimension remains because axis=0 only removes that dimension.

🚀 Application

expert

Using np.expand_dims and np.squeeze to reshape data

You have a 1D array of shape (10,). You want to reshape it to (10, 1) for a machine learning model input, then later revert it back to (10,) after prediction. Which sequence of operations correctly achieves this?

NumPy

import numpy as np
arr = np.arange(10)
# reshape for model input
arr_expanded = np.expand_dims(arr, axis=1)
# ... model prediction happens ...
# revert to original shape
arr_reverted = np.squeeze(arr_expanded, axis=1)
print(arr_reverted.shape)

A Use np.expand_dims(arr, axis=0) then np.squeeze(arr_expanded)

B Use np.expand_dims(arr, axis=0) then np.squeeze(arr_expanded, axis=0)

C Use np.expand_dims(arr, axis=1) then np.squeeze(arr_expanded)

D Use np.expand_dims(arr, axis=1) then np.squeeze(arr_expanded, axis=1)

💡 Hint

Match the axis used in expand_dims and squeeze to correctly revert shape.

Expanding at axis=1 changes shape from (10,) to (10, 1). Squeezing at axis=1 removes that added dimension, reverting to (10,). Using axis=0 in either step changes the shape incorrectly.