Available Models

Model	Video Model	CausalLM	AttentionModule	Parameters Quantization	Operation Bit Quantization
Gptj	❌	✅	✅	✅	✅
LucidTransformer	❌	✅	✅	✅	✅
Mixtral	✅	✅	✅	✅	✅
Opt	❌	✅	✅	✅	✅
Qwen2Moe	❌	✅	✅	✅	✅
Stablelm	❌	✅	✅	✅	✅
Cohere	❌	✅	✅	✅	✅
Arctic	❌	✅	✅	✅	✅
OpenELM	❌	✅	✅	✅	✅
Gemma	❌	✅	✅	✅	✅
GptNeoX	❌	✅	✅	✅	✅
Jetmoe	❌	✅	✅	✅	✅
Mamba	❌	✅	❌	✅	✅
MosaicMpt	❌	✅	✅	✅	✅
Palm	❌	✅	✅	✅	✅
Qwen1	❌	✅	✅	✅	✅
Roberta	❌	✅	✅	✅	✅
T5	❌	✅	✅	✅	✅
Dbrx	❌	✅	✅	✅	✅
Falcon	❌	✅	✅	✅	✅
Gpt2	❌	✅	✅	✅	✅
Grok1	❌	✅	✅	✅	✅
Llama	✅	✅	✅	✅	✅
Mistral	✅	✅	✅	✅	✅
Olmo	❌	✅	✅	✅	✅
Phi	❌	✅	✅	✅	✅
Phi 3	❌	✅	✅	✅	✅
Qwen2	❌	✅	✅	✅	✅
Rwkv	❌	✅	❌	✅	✅
Whisper	❌	✅	✅	✅	✅

you can also tell me the model you want in Flax/Jax version and ill try my best to build it ;)

More Models might have been added to ~HEAD but not mentioned here