Run Any HuggingFace Model on TPUs: A Beginner's Guide to TorchAX

What if you could run any HuggingFace model on TPUs — without rewriting a single line of model code? Here is what the end result looks like: from transformers import AutoModelForCausalLM , AutoTokenizer model = AutoModelForCausalLM . from_pretrained ( " google/gemma-3-1b-it " , torch_dtype = " bfloat16 " ) import torchax torchax . enable_globally () # Enable AFTER loading the model model . to ( " jax " ) # That's it. Now running on JAX. Five lines. Your PyTorch model is now executing on JAX — with access to TPUs, JIT compilation, and automatic parallelism across devices. In this tutorial, we will go from zero to building a working chatbot powered by a HuggingFace model running on JAX. Along the way, you will learn key JAX concepts, see real benchmarks, and understand why this approach exists. Why This Matters: The HuggingFace + JAX Problem In 2024, HuggingFace removed native JAX and TensorFlow support from its transformers library to focus development on PyTorch. This left thousands of