{"id":19464379,"url":"https://github.com/aryaaftab/rotary-embedding-tensorflow","last_synced_at":"2025-08-26T14:09:25.611Z","repository":{"id":57462710,"uuid":"400794024","full_name":"AryaAftab/rotary-embedding-tensorflow","owner":"AryaAftab","description":"Implementation of Rotary Embeddings, from the Roformer paper, in Tensorflow ","archived":false,"fork":false,"pushed_at":"2021-09-12T04:33:05.000Z","size":7,"stargazers_count":12,"open_issues_count":1,"forks_count":2,"subscribers_count":2,"default_branch":"master","last_synced_at":"2025-05-23T13:38:44.904Z","etag":null,"topics":["deeplearning","positional-encoding","tensorflow2"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/AryaAftab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2021-08-28T13:02:39.000Z","updated_at":"2024-11-02T08:32:11.000Z","dependencies_parsed_at":"2022-09-05T17:21:22.707Z","dependency_job_id":null,"html_url":"https://github.com/AryaAftab/rotary-embedding-tensorflow","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/AryaAftab/rotary-embedding-tensorflow","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AryaAftab%2Frotary-embedding-tensorflow","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AryaAftab%2Frotary-embedding-tensorflow/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AryaAftab%2Frotary-embedding-tensorflow/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AryaAftab%2Frotary-embedding-tensorflow/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/AryaAftab","download_url":"https://codeload.github.com/AryaAftab/rotary-embedding-tensorflow/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AryaAftab%2Frotary-embedding-tensorflow/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":261548735,"owners_count":23175495,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deeplearning","positional-encoding","tensorflow2"],"created_at":"2024-11-10T18:14:46.226Z","updated_at":"2025-06-23T20:07:15.204Z","avatar_url":"https://github.com/AryaAftab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Rotary Embeddings - Tensorflow\n\nA standalone library for adding \u003ca href=\"https://arxiv.org/abs/2104.09864\"\u003erotary embeddings\u003c/a\u003e to transformers in Tesnorflow, following its success as \u003ca href=\"https://blog.eleuther.ai/rotary-embeddings/\"\u003erelative positional encoding\u003c/a\u003e. Specifically it will make rotating information into any axis of a tensor easy and efficient, whether they be fixed positional or learned. This library will give you state of the art results for positional embedding, at little costs.\n\nMy gut also tells me there is something \u003ca href=\"https://www.nature.com/articles/s41593-021-00821-9\"\u003emore\u003c/a\u003e to rotations that can be exploited in artificial neural networks.\n\n## Note\nAn implemented version of Pytorch is available in this \u003ca href=\"https://github.com/lucidrains/rotary-embedding-torch\"\u003erepository\u003c/a\u003e.\n\nThis version is written by converting to the version of Pytorch. \n\nThe three functions of rearrange, irearrange and repeat have been written due to the incompatibility of the einops library with tensorflow 2.x.\n## Install\n\n```bash\n$ pip install rotary-embedding-tensorflow\n```\n\n## Usage\n\n```python\nimport tensorflow as tf\nfrom rotary_embedding_tensorflow import apply_rotary_emb, RotaryEmbedding\n\n# instantiate the positional embedding in your transformer and pass to all your attention layers\n\npos_emb = RotaryEmbedding(dim = 32)\n\n# generate the rotations\n\nfreqs = pos_emb(tf.range(1024), cache_key = 1024) # cache with a key that is the sequence length, so that it does not need to recompute\n\n# mock queries and keys\n\nq = tf.random.normal((1, 1024, 64)) # queries - (batch, seq len, dimension of head)\nk = tf.random.normal((1, 1024, 64)) # keys\n\n# apply the rotations to your queries and keys after the heads have been split out, but prior to the dot product and subsequent softmax (attention)\n\nfreqs = freqs[None, ...] # expand dimension for batch dimension\nq = apply_rotary_emb(freqs, q)\nk = apply_rotary_emb(freqs, k)\n\n# then do your attention with your queries (q) and keys (k)\n```\n\nIf you do all the steps above correctly, you should see a dramatic improvement during training\n\n## Axial Rotary Embeddings\n\nFor easy use of 2d axial relative positional embedding, ie. vision transformers\n\n```python\nimport tensorflow as tf\nfrom rotary_embedding_tensorflow import apply_rotary_emb, RotaryEmbedding, broadcat\n\npos_emb = RotaryEmbedding(\n    dim = 32,\n    freqs_for = 'pixel'\n)\n\n# queries and keys for frequencies to be rotated into\n\nq = tf.random.normal((1, 256, 256, 64))\nk = tf.random.normal((1, 256, 256, 64))\n\n# get frequencies for each axial\n# -1 to 1 has been shown to be a good choice for images and audio\n\nfreqs_h = pos_emb(tf.linspace(-1, 1, num = 256), cache_key = 256)\nfreqs_w = pos_emb(tf.linspace(-1, 1, num = 256), cache_key = 256)\n\n# concat the frequencies along each axial\n# broadcat function makes this easy without a bunch of expands\n\nfreqs = broadcat((freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim = -1)\n\n# rotate in frequencies\n\nq = apply_rotary_emb(freqs, q)\nk = apply_rotary_emb(freqs, k)\n```\n\n## Learned Rotations\n\nFor injecting learned rotations into a network. Experiments pending\n\nUpdate: doesn't seem to do anything -_-, will keep trying...\n\n```python\nimport tensorflow as tf\nfrom tensorflow.keras import layers\nfrom rotary_embedding_tensorflow import apply_learned_rotations\n\nx = tf.random.normal((1, 1024, 512))\n\n# you can only rotate in (dim // 2) values\n# ex. for 512, you can only rotate in 256 values\n\n# say you have two sets of learned rotations of 128 values each\n\nrots1 = layers.Dense(128)(x)\nrots2 = layers.Dense(128)(x)\n\n# you rotate in 256 (128 x 2) at first\n\nx = apply_learned_rotations(rots1, x, start_index = 0)\n\n# then you start at index 256 and rotate in the last (128 x 2)\n\nx = apply_learned_rotations(rots2, x, start_index = 256)\n\n# you could also concat the rotations together and pass it in all at once\n\nrots = tf.concat((rots1, rots2), axis = -1)\n\nx = apply_learned_rotations(rots, x)\n```\n\n## Citations\n\n```bibtex\n@misc{su2021roformer,\n    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding}, \n    author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},\n    year    = {2021},\n    eprint  = {2104.09864},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CL}\n}\n\n@misc{rotary-embedding-torch,\n    title   = {Rotary Embeddings - Pytorch}, \n    author  = {Phil Wang (lucidrains)},\n    year    = {2021},\n    url  = {https://github.com/lucidrains/rotary-embedding-torch},\n    publisher = {Github},\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Faryaaftab%2Frotary-embedding-tensorflow","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Faryaaftab%2Frotary-embedding-tensorflow","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Faryaaftab%2Frotary-embedding-tensorflow/lists"}