Does this library support 2D RoPE embeddings?

[logicchains]

The original RoPE paper suggests it can easily be extended to a 2D embedding (with two embedding dimensions rather than just one), does this library support that? If not, would it be difficult to add?

[lucidrains]

@logicchains yes it does, it is in the readme, under axial rotary embeddings

[lucidrains]

i also have an example here https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/rvt.py validated to converge well by some researchers on Eleuther