HF LLaMA Flax

[sanchit-gandhi]

Hey @ayaka14732! Super cool repo - thanks for working on this! With @vvvm23, we're working on adding the Flax LLaMA model to HF Transformers: https://github.com/huggingface/transformers/pull/24587 Just thought I'd let you know since it might be of interest to you, and potentially a model class that you can leverage in this repo for fast inference/training. It'll be a fully integrated version of the Flax LLaMA model in Transformers by the time of release.