OpenViSTAI · Agents · Blockchain · Crypto
Submit a project
Menu
AI Developer Tools & SDKs

Axial Attention

Axial Attention is a Python library implementing axial attention and related transformer components in PyTorch. Its `AxialAttention` class applies attention across configured dimensions of multidimensional tensors, and the package also provides axial positional embeddings and an image-transformer component.

Explore the practical guide ↓

Independent OpenViST project profile, based on linked sources.

Axial Attention — OpenViST project profile in AI Developer Tools & SDKs.
OpenViST project illustration. See the linked sources for product details.
Understand the project

Practical overview

When it fits

Use it while implementing a PyTorch model when attention over multiple dimensions of image or video data is needed and you want the library to handle the associated tensor reshaping and permutation.

Who it is for

  • PyTorch developers building image models with attention across spatial dimensions.
  • Researchers and developers working with video-shaped tensors and attention across multiple axes.
  • Model builders who want axial positional embeddings or an image-transformer component.

Documented uses

  • Add axial attention layers to image-processing models.
  • Configure attention over the axes of a video tensor.
  • Add axial positional embeddings to multidimensional image representations.
  • Build an image-transformer stack using `AxialImageTransformer`.
Before you choose

Requirements & limits

The README documents installation using `pip install axial_attention`. The repository identifies the project as MIT-licensed.

  • Set dim_index to identify where the embedding dimension sits in the input tensor.
  • Set num_dimensions to match the axes used for attention; the examples use two for images and three for video.
  • The component exposes head-count configuration, and sum_axial_out controls whether axis contributions are summed or applied sequentially.
  • The README shows positional embeddings with a specified spatial shape, such as (20, 20).

Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.

Source-based practical guide

Configure axial attention for an image tensor

Add a two-axis AxialAttention component to a PyTorch image tensor.

Before you start

  • Python environment with PyTorch
  • Install the package with pip install axial_attention

Steps

  1. Import torch and AxialAttention.

  2. Create an image tensor, as in the documented example with shape (1, 3, 256, 256).

  3. Instantiate AxialAttention with dim=3, dim_index=1, and num_dimensions=2; configure heads and dim_heads as appropriate for the example.

  4. Pass the tensor to the attention component.

Expected outcome

The example returns a tensor with shape `(1, 3, 256, 256)`.

Guide limitations

  • The example documents one tensor layout and configuration; adapt the embedding-dimension index and attention settings to your model's tensor.
Documentation 1 ↗
Source-based practical guide

Add axial positional embeddings to a multidimensional image representation

Apply the library's axial positional embedding and attention components to a channel-first image tensor.

Before you start

  • Python environment with PyTorch
  • Install the package with pip install axial_attention

Steps

  1. Import AxialAttention and AxialPositionalEmbedding.

  2. Create a tensor with the documented shape (1, 512, 20, 20).

  3. Configure AxialAttention with dim=512, heads=8, and dim_index=1.

  4. Create AxialPositionalEmbedding with dim=512 and shape=(20, 20), apply it to the tensor, then pass the result to attention.

Expected outcome

The example produces positionally embedded data and an attention output with shape `(1, 512, 20, 20)`.

Guide limitations

  • The documented positional-embedding example uses a 20-by-20 spatial shape and 512 embedding dimensions.
Documentation 1 ↗
Evidence and transparency

Sources & review

Official repository ↗
Source record details
Source availability
Active
Sources checked
16 October 2026
Original website title
lucidrains/axial-attention · Official homepage title · retained in its source language
How we review sourcesSuggest a correction