Rotary Position Embedding for Flax.

Implements the position embeddings from RoFormer: Enhanced Transformer with Rotary Position Embedding for Flax. rotary-embedding-torch was used as a reference implementation.

Features:

1D (for sequence models) and 2D axial (for ViT) rotary embeddings.
Learnable frequencies, including separate learnable frequencies per attention head.

Usage

Coming soon

Related Projects

TransformerEngine

A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating poin...

20 Sep 2022 1,482

RETRO-pytorch

Implementation of RETRO, Deepmind's Retrieval based Attention net, in Pytorch

15 Jan 2022 850

train-text-to-image-tpu-tutorial

20 Jan 2023 19

minGPT-flax

GPT implementation in Flax

30 Sep 2021 16

ECCV2022-RIFE

ECCV2022 - Real-Time Intermediate Flow Estimation for Video Frame Interpolation

12 Nov 2020 4,409

progen

Implementation and replication of ProGen, Language Modeling for Protein Generation, in Jax

09 Jun 2021 109

soundstorm-pytorch

Implementation of SoundStorm, Efficient Parallel Audio Generation from Google Deepmind, in Pytorch

17 May 2023 1,370

n-grammer-flax

Implementation of N-Grammer in Flax

04 Aug 2022 14

transformers

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

29 Oct 2018 132,459

lpips-j

Minimal JAX/Flax port of `lpips` supporting `vgg16`, with pre-trained weights stored in the 🤗 Hug...

08 Jul 2022 14

rotary-embedding-torch

Implementation of Rotary Embeddings, from the Roformer paper, in Pytorch

29 Jun 2021 535

flax

Flax is a neural network library for JAX that is designed for flexibility.

10 Jan 2020 5,584

fast-transformer-pytorch

Implementation of Fast Transformer in Pytorch

23 Aug 2021 171

performer-pytorch

An implementation of Performer, a linear attention-based transformer, in Pytorch

03 Oct 2020 1,084

TimeSformer-pytorch

Implementation of TimeSformer from Facebook AI, a pure attention-based solution for video classif...

11 Feb 2021 689

rope-flax