Fine-Tune Any HuggingFace Model like Gemma on TPUs with TorchAX

What if you could fine-tune any HuggingFace model on TPUs — using PyTorch code? Here is what the end result looks like: import torchax as tx import torchax.train # One function: forward → loss → gradients → optimizer update step_fn = tx . train . make_train_step ( model_fn , loss_fn , optimizer ) # Training loop for batch in dataloader : loss , params , opt_state = step_fn ( params , buffers , opt_state , batch , batch [ " labels " ]) Your PyTorch model. JAX's training primitives. Running on TPU