axial-positional-embedding

What it is and what it does

Axial Positional Embedding implements a specialized form of positional encoding designed for transformer networks working with multi-dimensional data. Instead of encoding position as a single sequence, it decomposes position information across multiple axes—useful for images, videos, or any structured tensor input. The package provides two main variants: a standard AxialPositionalEmbedding that combines embeddings from separate axes, and a ContinuousAxialPositionalEmbedding that uses MLPs for better extrapolation across unseen dimensions.

The package depends on torch for tensor operations and einops for flexible tensor reshaping. It targets developers building attention-based models on non-sequential data, where traditional sequence-based positional encodings are inefficient or inappropriate. The implementation is lightweight and integrates directly into PyTorch models by adding the embedding to token representations.

Use it for:

• Encoding position information in vision transformers processing image patches arranged in 2D grids
• Adding positional context to video transformers with temporal and spatial dimensions
• Improving transformer efficiency on multi-dimensional data by factorizing position encoding across axes
• Building language models that benefit from axial decomposition of sequence positions
• Extending transformers to handle variable-sized multi-dimensional inputs with better generalization

Worth the install?

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

Provides axial positional embeddings for transformer networks, enabling position encoding across multi-dimensional data like images, videos, or sequences using decomposed attention axes.

Yes, if you are building transformer models on multi-dimensional data (images, videos, or structured tensors) and want a lightweight, permissively licensed positional encoding layer. The aging maintenance status is a minor concern for a stable utility, but verify compatibility with your specific PyTorch version. No known vulnerabilities.

Install

axial-positional-embedding on PyPI

pip install axial-positional-embedding

uv add axial-positional-embedding

poetry add axial-positional-embedding

Installing axial-positional-embedding

Before you install

Low install friction with a pure Python wheel. Maintenance status is aging—last release was 535 days ago—but the package remains functional with torch and einops.

License in practice

MIT License permits free use, modification, and distribution with minimal restrictions, making it suitable for both open-source and commercial projects.

Quickstart

pip install axial-positional-embedding

import torch
from axial_positional_embedding import AxialPositionalEmbedding

pos_emb = AxialPositionalEmbedding(dim=512, axial_shape=(64, 64), axial_dims=(256, 256))
tokens = torch.randn(1, 1024, 512)
tokens = pos_emb(tokens) + tokens

Requires torch and einops as runtime dependencies.

Verify before relying

• Whether the package works with recent PyTorch versions beyond current support
• Performance characteristics or benchmarks compared to standard positional encoding
• Active community support or responsiveness to issues given aging maintenance status

Package facts

Evidence: axial_positional_embedding-0.3.12-py3-none-any.whl

Keywords: artificial intelligence, deep learning, positional embedding

Tags

More Artificial Intelligence packages

Further reading