numpy.squeeze() in Python

The numpy.squeeze() is a useful Python function, which is utilized for the removal of single-dimensional elements from the shape of a NumPy array. It comes in very handy when you have to discard redundant dimensions (like a dimension with size 1) after operations that introduce extra dimensions.

Basic usage of numpy.squeeze()

import numpy as np

in_arr = np.array([[[2, 2, 2], [2, 2, 2]]])
 
print ("Input array : ", in_arr) 
print("Shape of input array : ", in_arr.shape)  

out_arr = np.squeeze(in_arr) 

print ("output squeezed array : ", out_arr)
print("Shape of output array : ", out_arr.shape) 

Output :

Input array : [[[2 2 2]
[2 2 2]]]
Shape of input array : (1, 2, 3)
output squeezed array : [[2 2 2]
[2 2 2]]
Shape of output array : (2, 3)

Explanation:

• Input: A 3D array of shape (1, 2, 3)—the first dimension has size 1.
• Output: The numpy.squeeze() function removes the first dimension of size 1, resulting in a 2D array of shape (2, 3).

Syntax of numpy.squeeze() in Python

numpy.squeeze(arr, axis=None )

Parameters:

• arr: Input array.
• axis: Selects a subset of the single-dimensional entries in the shape. If an axis is selected with shape entry greater than one, an error is raised.

Return Type: The function returns a new array, which is a view of the input array with the single-dimensional entries removed from its shape

Example 1: Using the axis parameter to squeeze a specific dimension

import numpy as np
in_arr = geek.arange(9).reshape(1, 3, 3) 

print ("Input array : ", in_arr)  
out_arr = np.squeeze(in_arr, axis = 0) 

print ("output array : ", out_arr)  
print("The shapes of Input and Output array : ") 
print(in_arr.shape, out_arr.shape)

Output :

Input array : [[[0 1 2]
[3 4 5]
[6 7 8]]]
output array : [[0 1 2]
[3 4 5]
[6 7 8]]
The shapes of Input and Output array :
(1, 3, 3) (3, 3)

Explanation: The input array has shape (1, 3, 3). By specifying axis = 0, we remove the first dimension, resulting in an output array with shape (3, 3).

Example 3: Error when trying to squeeze a non-singleton dimension

import numpy as np

in_arr = np.arange(9).reshape(1, 3, 3)
print("Input array: ", in_arr)  

# Trying to squeeze the axis that doesn't have size 1
try:
    out_arr = np.squeeze(in_arr, axis=1)  
except ValueError as e:
    print("Error: ", e

Output :

Input array: [[[0 1 2]
[3 4 5]
[6 7 8]]]
Error: cannot select an axis to squeeze out which has size not equal to one

Explanation: The input array has a shape of (1, 3, 3). The attempt to squeeze axis = 1, which corresponds to the second dimension with size 3, results in a ValueError, as the dimension is not of size 1.