Of course! In Python, the function you're looking for is unsqueeze. It's a core operation in libraries like PyTorch and NumPy, although their names and implementations differ slightly.

The core idea of unsqueeze is to add a dimension of size 1 to a tensor or array. This is extremely useful for making tensors compatible for operations that require specific shapes, especially in machine learning.

Let's break it down for both PyTorch and NumPy.

PyTorch (torch.unsqueeze)

In PyTorch, the function is torch.unsqueeze(). It's a very common operation when preparing data for models.

What it does

It adds a new dimension at a specified position. The size of this new dimension is always 1.

Syntax

torch.unsqueeze(input, dim) 
# or the more readable version:
input.unsqueeze(dim)

• input: The PyTorch tensor you want to modify.
• dim (or index): The index where you want to insert the new dimension.

Key Concept: Broadcasting

The main reason for using unsqueeze is broadcasting. Broadcasting is a set of rules that allows PyTorch to perform operations on tensors of different shapes. To add a dimension, the original dimension sizes must be equal, or one of them must be 1.

Example: Adding a Channel Dimension Imagine you have a grayscale image and you want to use a model that expects a color image (with 3 channels: Red, Green, Blue).

• A grayscale image might have the shape [height, width], e.g., [28, 28].
• A color image has the shape [channels, height, width], e.g., [3, 28, 28].

To make the grayscale image compatible, you need to add a channel dimension.

import torch
# A grayscale image (batch of 1 image)
# Shape: [1, height, width] -> [1, 28, 28]
grayscale_image = torch.randn(1, 28, 28) 
print(f"Original shape: {grayscale_image.shape}")
# Add a channel dimension at the beginning (dim=0)
# The new shape will be [1, 1, 28, 28]
image_with_channel = grayscale_image.unsqueeze(dim=0)
print(f"Shape after unsqueeze(0): {image_with_channel.shape}")
# Add a channel dimension between batch and height (dim=1)
# This is the most common operation for converting grayscale to "pseudo-color"
# The new shape will be [1, 1, 28, 28]
image_with_channel_dim1 = grayscale_image.unsqueeze(dim=1)
print(f"Shape after unsqueeze(1): {image_with_channel_dim1.shape}")

Output:

Original shape: torch.Size([1, 28, 28])
Shape after unsqueeze(0): torch.Size([1, 1, 28, 28])
Shape after unsqueeze(1): torch.Size([1, 1, 28, 28])

NumPy (numpy.expand_dims)

NumPy does not have a function named unsqueeze. Its equivalent is numpy.expand_dims(). The concept is identical.

What it does

It adds a new axis (dimension) to an ndarray, with a size of 1, at the specified position.

Syntax

numpy.expand_dims(a, axis)

• a: The input NumPy array.
• axis: The position in the expanded axes where the new axis (dimension) is placed.

Example: The same Grayscale Image Scenario

Let's replicate the previous example using NumPy.

import numpy as np
# A grayscale image (batch of 1 image)
# Shape: [1, height, width] -> [1, 28, 28]
grayscale_image_np = np.random.rand(1, 28, 28)
print(f"Original shape: {grayscale_image_np.shape}")
# Add a channel dimension at the beginning (axis=0)
# The new shape will be [1, 1, 28, 28]
image_with_channel_np = np.expand_dims(grayscale_image_np, axis=0)
print(f"Shape after expand_dims(axis=0): {image_with_channel_np.shape}")
# Add a channel dimension between batch and height (axis=1)
# The new shape will be [1, 1, 28, 28]
image_with_channel_dim1_np = np.expand_dims(grayscale_image_np, axis=1)
print(f"Shape after expand_dims(axis=1): {image_with_channel_dim1_np.shape}")

Output:

Original shape: (1, 28, 28)
Shape after expand_dims(axis=0): (1, 1, 28, 28)
Shape after expand_dims(axis=1): (1, 1, 28, 28)

The Reverse Operation: squeeze

The opposite of unsqueeze is squeeze. It removes dimensions of size 1.

• PyTorch: torch.squeeze(input) or input.squeeze(). If dim is specified, it only removes that dimension if its size is 1. If no dim is given, it removes all dimensions of size 1.
• NumPy: numpy.squeeze(a). Same behavior as PyTorch.

Example:

import torch
import numpy as np
# A tensor with several dimensions of size 1
# Shape: [1, 1, 3, 1, 5]
tensor = torch.randn(1, 1, 3, 1, 5)
print(f"Original shape: {tensor.shape}")
# Remove all dimensions of size 1
squeezed_tensor = torch.squeeze(tensor)
print(f"Squeezed shape (all dims of size 1 removed): {squeezed_tensor.shape}") # -> [3, 5]
# Remove a specific dimension of size 1
squeezed_dim1 = torch.squeeze(tensor, dim=1)
print(f"Squeezed shape (dim=1 removed): {squeezed_dim1.shape}") # -> [1, 3, 1, 5]

Summary Table

When to Use unsqueeze / expand_dims

1. Machine Learning Data Preprocessing: The most common use case. Models often expect inputs in a specific format (e.g., [batch_size, channels, height, width]). Your raw data might not have all these dimensions.
2. Vector Operations: You might have a feature vector of shape [features] and need to perform an operation with a matrix of shape [batch_size, features]. You can unsqueeze the vector to make it [1, features] so it can be broadcast across the batch.
3. Creating Single-Item Batches: If you have one item and want to run it through a model that expects a batch, you can unsqueeze the batch dimension (usually dim=0) to go from [height, width] to [1, height, width].