Software and platforms
Machine learning for JAX
JAX is used to train neural networks, run large-scale research experiments and build scientific models that need exact gradients. You write NumPy-style functions, and JAX differentiates them with grad, compiles them with jit and batches them with vmap on CPU, GPU or TPU.
What JAX is used for in machine learning
Training and fine-tuning large language models
JAX compiles training steps with XLA and shards them across GPUs and TPUs, which suits very large models. Model weights are plain arrays, so fine-tuning code stays close to the maths.
Real exampleGemma library (Google DeepMind)Google DeepMind's gemma repository is a JAX library to use and fine-tune its open-weights Gemma language models.
Scientific models such as protein structure prediction
Research groups use JAX for large scientific networks because it differentiates and compiles numerical code efficiently on accelerators.
Real exampleAlphaFold 2 open-source codeDeepMind's open-source AlphaFold 2 inference code lists JAX and the Haiku neural network library among the libraries it is built with.
Differentiable physics and reinforcement learning
A simulator written in JAX can run thousands of environments in parallel on one accelerator and return gradients through the physics. That makes training control policies much faster than stepping a separate simulator on a CPU.
Real exampleBraxBrax is a fully differentiable physics engine written in JAX, with learning algorithms such as PPO and SAC for training agents.
Probabilistic and Bayesian modelling
Bayesian models estimate uncertainty in their parameters by sampling, which needs many gradient evaluations. JAX's grad and jit make samplers such as NUTS fast enough for practical models.
Real exampleNumPyroNumPyro is a probabilistic programming library powered by JAX for autograd and JIT compilation to GPU, TPU or CPU.
Can you learn machine learning with JAX?
Yes, if you are comfortable with Python and NumPy. JAX asks you to write the maths of training directly, so derivatives, matrix shapes and gradient descent come first, then a hand-written training loop, then Flax and Optax.
- Python and NumPy arraysJAX mirrors the NumPy API, so fluency with arrays, broadcasting and pure functions makes everything else easier.Learn it in Data Tools
- Derivatives and linear algebrajax.grad returns the gradient of the function you write, so you need to know what a gradient and a matrix product are.Learn it in Maths
- Models, losses and evaluationLearn regression, classification, loss functions and held-out evaluation, so you know what your JAX training loop should optimise.Learn it in Machine Learning
- Neural networks and optimisersUnderstand layers, backpropagation and optimisers such as Adam before using Flax modules and Optax.Learn it in Deep Learning
- JAX transformationsWrite one training loop with grad, jit and vmap by hand, then learn random keys, pytrees and how jit traces functions.
Moving into machine learning work with JAX
If you already work in the field
If you already write PyTorch or NumPy code, the main change is the functional style: pure functions, explicit random keys and parameters passed as data. Port one small PyTorch model to JAX and Flax to learn the differences.
If you come from data science or ML
If you come from scientific computing, physics or applied maths, you already think in arrays and derivatives, which is how JAX works. Add model training practice and evaluation, then use JAX for one simulation or fitting problem from your own field.
Roles that do this workResearch engineerMachine learning researcherMachine learning engineerComputational scientist
Tools you will use
| Tool | What you use it for |
|---|---|
| JAX | Write array code and get automatic gradients, compilation and vectorisation for training loops. |
| Flax | Define neural-network layers and models on top of JAX for research and production training. |
| Optax | Apply optimisers such as Adam and SGD, learning-rate schedules and gradient clipping to JAX parameters. |
| Equinox | Build models as simple Python classes that work directly with JAX transformations. |
Public data to practise on
- Concrete Compressive Strength (UCI): Mix ingredients and curing age for concrete samples; regress compressive strength with a small network.
- CIFAR-10: Small labelled colour images in ten classes; train a convolutional network in Flax and compare optimisers.
Your first project
Download the Concrete Compressive Strength data, predict strength with a small multilayer perceptron trained using jax.grad and Optax, and compare it with a linear regression baseline using mean absolute error on a held-out split.
Common questions
Is JAX better than PyTorch?
Neither is better in general. JAX suits functional code, compiled training on TPUs and research that needs custom transformations, while PyTorch has a larger ecosystem and more tutorials.
Is JAX hard to learn?
The API is small and close to NumPy. The harder parts are its rules: functions must be pure, random numbers need explicit keys, and jit-compiled code cannot use Python control flow on traced values.
Do I need a TPU to use JAX?
No. JAX runs on CPU and GPU as well, and the same code runs on TPU when one is available, for example in Google Colab.
What is the difference between JAX, Flax and Optax?
JAX provides arrays, gradients and compilation. Flax adds neural network layers and model structure, and Optax adds optimisers and learning-rate schedules.
Courses for these skills
- Maths for Machine LearningWork with norms, dot products and vector geometry.26 published lessons · all planned lessons open
- Deep Learning with PyTorchStudy sequence models and neural optimisation.Syllabus available · lessons open as produced
- Generative AI: LLMs, RAG and AgentsConnect embeddings and retrieval to language-model applications.Syllabus available · lessons open as produced
New to programming or maths? Start with Python, or follow the machine learning path in order. Published lessons are free to read; compare Free and Pro.