poly-attention 0.0.2

PolyAttention

Poly Attention (wip)

Implementation of Poly-Attention, a general scheme for higher-order self-attention

Install

$ pip install poly-attention

Usage

import torch
from poly_attention import PolyAttention

attn = PolyAttention(
    dim = 512,
    heads = 8,
    dim_head = 64,
    causal = False
)

x = torch.randn(1, 1024, 512)
out = attn(x) # (1, 1024, 512)

Citations

@inproceedings{chakrabarti2026poly,
    title   = {Poly-attention: a general scheme for higher-order self-attention},
    author  = {Chakrabarti, Sayak and Pitassi, Toniann and Alman, Josh},
    booktitle = {International Conference on Learning Representations (ICLR)},
    year    = {2026}
}