unfoldNd

What it is and what it does

unfoldNd generalizes PyTorch's im2col (unfold) and col2im (fold) operations to tensors with more than 4 dimensions. PyTorch's native torch.nn.functional.unfold and torch.nn.Unfold only work on 4D batched image tensors; this package extends that functionality to 3D, 4D, and 5D inputs using a numerical trick based on one-hot kernels and group convolutions.

The package is useful when you need to express convolutions as matrix-matrix multiplications on higher-dimensional data—a perspective that enables certain optimization techniques and automatic differentiation patterns. It trades some additional peak memory usage for speed gains in both forward and backward passes compared to native PyTorch operations. The package also exposes transpose convolution unfolding (unfoldTransposeNd) and fold operations (foldNd), though fold is tested but not benchmarked.

Use it for:

• Implementing second-order optimization methods like KFAC that approximate the Fisher matrix using Kronecker factors on volumetric (3D) convolutions.
• Expressing 3D or 5D convolutions as explicit matrix multiplications for custom automatic differentiation or layer implementations.
• Unfolding inputs for transpose convolutions in higher dimensions where PyTorch provides no native support.
• Research on convolution-as-linear-layer perspectives in deep learning frameworks like BackPACK.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Extends PyTorch's unfold and fold operations to 3D, 4D, and 5D tensors, implementing im2col and col2im for higher-dimensional inputs that PyTorch's native functions do not support.

Yes, if you work with 3D, 4D, or 5D convolutions and need the im2col perspective for optimization or custom AD. The package is production-stable, actively maintained, has no known vulnerabilities, and offers measurable speed improvements over workarounds. Install only if you actually need higher-dimensional unfold; it's a specialized tool, not a general PyTorch enhancement.

Install

unfoldnd on PyPI

pip install unfoldnd

uv add unfoldnd

poetry add unfoldnd

Installing unfoldNd

Before you install

Low friction installation with a pure Python wheel. The package is actively maintained with a recent release (2024-12-30) and depends only on packaging, torch, and numpy—all standard in PyTorch workflows.

License in practice

MIT license permits unrestricted use, modification, and distribution in both open-source and commercial projects with minimal obligations.

Quickstart

pip install unfoldNd

import unfoldNd
import torch

# Unfold a 5D tensor (batch, channels, depth, height, width)
x = torch.randn(2, 3, 4, 5, 6)
unfolded = unfoldNd.unfoldNd(x, kernel_size=3, padding=1)
# Returns unfolded view suitable for convolution-as-matrix-multiplication

Requires PyTorch and NumPy installed; intended for use with higher-dimensional tensor operations in convolution-based workflows.

Verify before relying

• Whether fold operations (col2im) have equivalent performance to unfold operations, as the description notes fold is not benchmarked.
• Memory overhead characteristics in production settings beyond the documented benchmark configurations.
• Compatibility with recent PyTorch versions beyond 3.10 support stated in classifiers.

Package facts

Evidence: unfoldNd-0.2.3-py3-none-any.whl

Keywords: PyTorch, im2col, col2im, unfold, fold

Tags

More Artificial Intelligence packages