Finding Minimum and Maximum Values in PyTorch Tensors with torch.Tensor.aminmax()

Purpose

Computes the minimum and maximum values along a specified dimension (or across the entire tensor if no dimension is given).

Syntax

torch.Tensor.aminmax(*, dim=None, keepdim=False) -> (Tensor min, Tensor max)

Parameters

keepdim (bool, optional): If True, the output tensors min and max will have the same dimensions as the input tensor, with a dimension of size 1 inserted at the specified dim. This is useful for further calculations that require the same dimensionality. Defaults to False.
dim (Optional[int]): The dimension over which to find the minimum and maximum values. If None (default), the operation is performed across the entire tensor.

Return Value

A named tuple (min, max) containing two tensors:
- min: A tensor containing the minimum values along the specified dimension.
- max: A tensor containing the maximum values along the specified dimension.

Example

import torch

tensor = torch.tensor([[3, 1, 4], [2, 5, 0]])

# Find minimum and maximum across the entire tensor (dim=None)
min_max = tensor.aminmax()
print(min_max)  # (tensor(0), tensor(5))

# Find minimum and maximum along the first dimension (rows)
min_max = tensor.aminmax(dim=0)
print(min_max)  # (tensor(2, 1, 0), tensor(3, 5, 4))

# Find minimum and maximum along the second dimension (columns) with keepdim=True
min_max = tensor.aminmax(dim=1, keepdim=True)
print(min_max)  # (tensor([[3], [2]]), tensor([[4], [5]]))

Key Points

It supports gradients, meaning you can use it in autograd computations.
aminmax() can handle tensors with different data types, including floating-point numbers and integers.
keepdim helps maintain the original dimensionality of the tensor for further calculations.
The dim parameter allows you to control which dimension to operate on.
aminmax() is a convenient way to find minimum and maximum values along a specific dimension of a tensor.

NaN (Not a Number) values in the tensor are propagated to the output if at least one value is NaN.
If the input tensor has a dimension with size 0, aminmax() will raise a RuntimeError.

Min-Max Normalization

This code normalizes a tensor data to a range between 0 and 1 using min-max normalization:

import torch

data = torch.tensor([[-1, 3, 2], [4, -2, 1]])

# Find minimum and maximum values across the entire tensor
min_val, max_val = data.aminmax()

# Normalize each element using broadcasting
normalized_data = (data - min_val) / (max_val - min_val)

print(normalized_data)

Finding Minimum/Maximum Values in Specific Locations

This code finds the minimum and maximum values along the second dimension (columns) and then retrieves their indices:

import torch

data = torch.tensor([[3, 1, 4], [2, 5, 0]])

# Find min/max and their indices along second dimension (columns)
min_values, min_indices = data.amin(dim=1, keepdim=True)
max_values, max_indices = data.amax(dim=1, keepdim=True)

print(min_values)  # tensor([[3], [2]]) (minimum values)
print(min_indices)  # tensor([[0], [0]]) (indices of minimum values)
print(max_values)  # tensor([[4], [5]]) (maximum values)
print(max_indices)  # tensor([[2], [1]]) (indices of maximum values)

Handling Tensors with NaN Values

This code demonstrates how aminmax() propagates NaNs:

import torch

data = torch.tensor([[float('nan'), 3, 2], [4, -2, float('nan')]])

# Find minimum and maximum
min_val, max_val = data.aminmax()

print(min_val)  # tensor(nan)
print(max_val)  # tensor(nan) (because at least one NaN exists)

Using torch.min() and torch.max()

These functions can be used separately to achieve the same result as aminmax(). However, they require separate calls for minimum and maximum values, and managing dimensions might be less intuitive.

import torch

tensor = torch.tensor([[3, 1, 4], [2, 5, 0]])

# Find minimum across the entire tensor
min_val = torch.min(tensor)  # Returns a single value

# Find maximum along the first dimension (rows)
max_val, _ = torch.max(tensor, dim=0)  # Ignores the second output (indices) 

# Combine them if necessary
combined_min_max = (min_val, max_val)

Manual Looping

For simple cases, or if you need more control over the calculations, you can use a loop to iterate through the tensor and track the minimum and maximum values yourself. However, this approach is less efficient for larger tensors.

import torch

tensor = torch.tensor([[3, 1, 4], [2, 5, 0]])

min_val = float('inf')
max_val = float('-inf')

for row in tensor:
    min_val = min(min_val, row.min())
    max_val = max(max_val, row.max())

combined_min_max = (min_val, max_val)

NumPy Integration (if applicable)

If you're working with tensors that can be converted to NumPy arrays, you can leverage NumPy's min() and max() functions for potentially faster performance on CPUs. However, this approach involves data transfer between PyTorch and NumPy, which might introduce overhead.

import torch
import numpy as np

tensor = torch.tensor([[3, 1, 4], [2, 5, 0]])

# Convert to NumPy array
numpy_array = tensor.numpy()

# Find min/max using NumPy
min_val = np.min(numpy_array)
max_val = np.max(numpy_array)

combined_min_max = (min_val, max_val)

Choosing the best alternative depends on factors like:

Specific needs
If you need more control over calculations or are using CPUs with NumPy integration, consider the alternatives.
Performance
For larger tensors on GPUs, torch.Tensor.aminmax() is likely the most efficient choice.
Readability and maintainability
torch.Tensor.aminmax() is generally more concise and easier to maintain.