Llama and Falcon in JAX

[rcrowe-google]

So far I've found at least 4 ports of Llama to JAX:

• Orig Llama: https://huggingface.co/heegyu/llama-7b-easylm
• Orig Llama: https://github.com/Sea-Snell/JAX_llama
• Llama2: https://github.com/ayaka14732/llama-2-jax
• Llama2: https://huggingface.co/erfanzar/Llama-2-jax

I'm wondering, are there are others that folks have done?

Also, I've looked for ports of Falcon to JAX and so far I haven't found any. Do you know of any?

[davisyoshida]

Here's my haiku port of Llama: https://github.com/davisyoshida/llama-haiku

Haiku has perished though 💔

[rcrowe-google]

Haiku has perished though 💔

Thanks for the link Davis, and sorry about Haiku.

[patrick-kidger]

Looks like you can find this in Levanter/Equinox: stanford-crfm/levanter#270

[rcrowe-google]

Thanks Patrick! Please let me know if you find a port of Falcon, I still haven't found one.

[ayaka14732]

Greetings from the author of llama-2-jax 👋🏻

[rcrowe-google]

Sorry, I forgot to include this one - https://github.com/openlm-research/open_llama

[defdet]

A bit late to the party, but there is falcon in jax now https://github.com/defdet/falcon-jax