Demo: JAX, Flax and Gemma
summary
Claims that JAX is the core engine behind Gemma and that it yields measurable speed and cost wins—citing a Vision Transformer study where JAX runs faster and cheaper than PyTorch—and supports training at scale (Gemma development trains on up to six trillion tokens). Uses Flax, Orbax and Optax for model code, demonstrates LoRA fine-tuning with Flax on a responsibly curated dataset, and shows loading the Gemma 2B instruction‑tuned model via Kaggle in four short steps: download checkpoints, init tokenizer, setup transformer, start sampling.
AI-written summary · about
related talks
09mUsing and fine-tuning Gemma 3
09mOpen models and open compression
26mTraining Large Deep Learning Models: PyTorch vs JAX
18mLive demo: Gemma 3 on Jetson Orin Nano
39mFlutter Project with Gemma: A State-of-the-Art Open Model
11mIntro to JAX: Accelerating Machine Learning research
14mGemma on mobile and web. Best and worst practices
44mBuilding amazing Front-End Project with Gemma: A State-of-the-Art Open Model