The original RoPE paper suggests that the Rotary Position Embedding it describes can easily be extended to two or more dimensions: 3.2.2 in https://arxiv.org/abs/2104.09864 . I'm trying to find a PyTorch implementation of this and I can't; I can only find implementations for a 1D embedding. Is anyone aware of any 2D implementations in PyTorch (or another framework, which I could translate to PyTorch)?

  • 1
    $\begingroup$ Not an expert in rotary embeddings or vision transformers but, after Googling, I bumbped into the term "axial rotary embeddings" and into these two pytorch libraries that support them: rotary-embedding-torch, axial-positional-embedding. Is it what you were looking for? $\endgroup$
    – noe
    Oct 9, 2023 at 14:17


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.