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 JAX users — especially