llama-models icon indicating copy to clipboard operation
llama-models copied to clipboard
verified publisher icon meta-llama

JAX Llama 4 implementation

Open SamuelMarks opened this issue 8 months ago • 0 comments

WiP; creating PR now to get some feedback.

Once ready would be interesting to see benchmarks on your same hardware with JAX as compared with PyTorch.

Hopefully you having made PyTorch doesn't make a JAX contrib unwanted!

SamuelMarks avatar Apr 06 '25 18:04 SamuelMarks