Registry / ai-ml / axial-positional-embedding

axial-positional-embedding

JSON →
library0.3.12pypypi✓ verified 85d ago

Implementation of Axial Positional Embedding, as described in 'Attention is All You Need' and 'Axial Attention in Transformers'. This library provides a simple way to add positional encodings to transformer models in an axial (factorized) manner, reducing the number of parameters. Current version 0.3.12, supports Python >=3.8.

pip install axial-positional-embedding
INSTALL
IMPORT
SIG · AXIAL-POSITIONAL-E
A
axial-positional-embedding
ai-mlpythonv0.3.12
Install
65.8s avg
Import
5400ms
Disk
4787MB
Pass rate
4/ 10
Env Coverage4 / 10
glibc
3.93.13
musl
3.93.13
Install & Compatibility
Where this runs
tested against v0.3.12 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
glibc
py 3.10
✕ build_error
✓ 73.6s
py 3.11
✕ build_error
✓ 68s
py 3.12
✕ build_error
✓ 63.7s
py 3.13
✕ build_error
✓ 57.7s
py 3.9
✕ build_error
✕ timeout
4787MB installed
● package 4787MB
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

AxialPositionalEmbedding
from axial_positional_embedding import AxialPositionalEmbedding
from axial_positional_embedding.axial_positional_embedding import AxialPositionalEmbedding
Common mistake: deep import path; top-level import is correct.

Basic usage: create an axial positional embedding layer for 2D spatial data and apply it to an input tensor.

import torch from axial_positional_embedding import AxialPositionalEmbedding # Example: create an embedding layer for a 2D image (height, width) dim = 128 # For an input of shape (batch, seq_len, dim) or (batch, height, width, dim)? # Typically, AxialPositionalEmbedding expects shape (batch, height, width, dim) for 2D axial. # But the library is flexible; let's assume a 2D spatial input. batch = 2 height = 16 width = 16 pos_emb = AxialPositionalEmbedding(dim=dim, shape=(height, width)) # Generate dummy input (batch, height, width, dim) x = torch.randn(batch, height, width, dim) # Apply positional embedding out = pos_emb(x) print(out.shape) # Expected: (batch, height, width, dim)
Debug
Known issues
gotchaInput tensor shape: The library expects input shape (batch, height, width, channels) for 2D axial, not (batch, channels, height, width). If your data is in channel-first format, you must permute before passing.
fix
If using PyTorch's usual image format (NCHW), permute to (batch, height, width, channels) before applying: x = x.permute(0, 2, 3, 1)
affects: all
gotchaDimension mismatch: The 'dim' parameter must match the channel dimension of the input tensor. Common mistake: setting dim=128 but input has 64 channels.
fix
Ensure dim equals the last dimension of the input tensor after permutation.
affects: all
deprecatedThe library does not actively deprecate features, but the original 'axial-positional-embedding' may see reduced updates. For production, consider using Hugging Face Transformers' built-in axial positional embedding.
fix
Check the GitHub repo for latest changes; if lacking, evaluate alternatives.
affects: 0.3.x
Errors
Common errors & fixes
ImportError: cannot import name 'AxialPositionalEmbedding'
Import path incorrect; likely using a wrong module name.
fix
Use: from axial_positional_embedding import AxialPositionalEmbedding
RuntimeError: The expanded size of the tensor must match the existing size
Input shape mismatch: the positional embedding expects a specific shape (height, width) and may broadcast incorrectly if dimensions don't align.
fix
Ensure input tensor has shape (batch, height, width, dim) after permutation.
AssertionError: dim must be divisible by 2?
Some implementations require even dimension for sin/cos positional encoding; this library currently does not enforce but version may vary.
fix
If you encounter this error, ensure dim is even.
Upgrade
Version history
0.3.12latest on PyPI · released Feb 25, 2025
Audit
Dependencies
torchrequiredCore dependency; library requires PyTorch tensors and modules.
einopsrequiredUsed for tensor rearranging in positional embedding.
Agent activity
13 hits · last 30 days
node
12
OpenAI (training)
1
Resources
axial-positional-embedding — pip install axial-positional-embedding · libregistry