feat: support ROPE - Githubissues

patil-suraj / vit-vqgan

JAX implementation ViT-VQGAN

MIT License

77 stars 11 forks source link

feat: support ROPE #23

Closed borisdayma closed 2 years ago

borisdayma commented 2 years ago

Support Rotary Embeddings.

Notes:

I try to avoid position_ids and indexing if not necessary
This is applied in Attention Blocks
This uses the dimension per head (not total)