ML frameworks

Libraries for training and running neural networks: arrays on GPUs and TPUs, automatic gradients, and a compiler that turns Python model code into fast kernels.

1 technology on this page ↓

How it works

Every framework rests on three pieces. The first is an array type, called a tensor, that lives on a CPU, GPU or TPU, with hundreds of operations (matrix multiply, convolution, softmax) implemented as kernels for each kind of hardware. The second is automatic differentiation. Training means computing a loss, then the gradient of the loss with respect to every weight, then nudging each weight against its gradient; nobody writes those derivatives by hand. Reverse-mode autodiff, or backpropagation, records the operations of the forward pass and replays them backwards to get all gradients in one pass. The third is a way to make Python fast enough. Python itself is slow, so frameworks capture the operations as a graph and run it in C++, CUDA or a compiler. They differ in when they capture it. PyTorch runs eagerly: each line executes as Python reaches it, and autograd records the graph as you go, “recreated from scratch at every iteration”, which allows any Python control flow. Since PyTorch 2.0 (March 2023), torch.compile can capture and compile that code for speed. TensorFlow 2 is also eager by default, and tf.function traces a function into a graph that runs without Python and can be saved and served. JAX traces pure functions and compiles them with XLA, and treats differentiation, compilation and batching as transformations you can combine. On top of the core sit layers, optimizers, data loaders, checkpointing and ways to split a model over many accelerators.

Think of it as

A spreadsheet that remembers every formula. Change a cell and it recalculates the total (the forward pass), and because it knows each formula, it can also tell you how much the total would move if you nudged any input (the gradient). A framework is that spreadsheet for billions of cells, with a calculator chip for each column.

Reach for it when

You train or fine-tune your own models, from a small classifier to a large transformer
You do research on new architectures or training methods and need to change the model code itself
You need custom numerical code on GPUs or TPUs that uses gradients, such as physics simulation or optimization, which is common with JAX
You run an open-weight model inside your own Python code rather than through a serving runtime

Think twice when

You call a hosted model API such as Claude or Gemini: the provider runs the framework, and your code sends text over HTTP
You only need to serve an open-weight LLM: a runtime such as llama.cpp, Ollama or vLLM loads the weights and serves an API with no training code
The data is a small table and a classic model will do: gradient-boosted trees or linear regression are often quicker to train and easier to explain than a neural network

A real example

A team fine-tunes a 300-million-parameter image classifier on 50,000 labelled product photos. In PyTorch, the training loop loads a batch of 256 images onto a GPU, runs the model forward, computes cross-entropy loss, calls loss.backward() to fill in every weight’s gradient, and lets an AdamW optimizer update the weights; wrapping the model in torch.compile fuses operations and shortens each step. The same job in JAX would write the loss as a pure function, take jax.grad of it, compile the update with jax.jit, and update the weights with an Optax optimizer. Either way the result is a checkpoint of weights that a serving system loads later.

Technologies

Specific products, and what each does differently.

JAXOpen sourceNumPy with function transformations: grad differentiates, jit compiles with XLA, vmap vectorizes, and the same code runs on CPUs, GPUs and TPUs across thousands of devices.


JAX	Open source	Google	Apache 2.0	2026-10-24

What each provider calls it

AWS: AWS Deep Learning Containers: prebuilt images with PyTorch or TensorFlow (and vLLM, SGLang and Ray) for training and inference on EC2, EKS and Amazon SageMaker AI
Google Cloud: Cloud TPU VMs, with guides for running JAX and PyTorch on TPUs; Google trains Gemini with JAX and Pathways
Azure: Azure Machine Learning curated environments: Docker images with the Python packages for common machine-learning frameworks, ready to use in training jobs
Open source & others: PyTorch (PyTorch Foundation, Linux Foundation), JAX with Flax and Optax (Google), TensorFlow (Google), and Keras 3, which runs on JAX, TensorFlow or PyTorch

When the graph is captured#

All three record the operations so they can compute gradients and run fast; they differ in when they turn Python into a graph.

	PyTorch	TensorFlow 2	JAX
Default mode	Eager: each operation runs when Python reaches it	Eager: operations run one by one and return to Python	Eager for single operations; real work goes through `jax.jit`
Gradients	Autograd records a graph during the forward pass; `loss.backward()` walks it	`tf.GradientTape` records operations inside its block	`jax.grad(f)` returns a new function that computes the gradient
Getting a compiled graph	`torch.compile` (since 2.0): TorchDynamo captures, TorchInductor generates kernels, using Triton on GPUs	`tf.function` traces into a graph; AutoGraph converts `if` and loops into graph ops	`jax.jit` traces a pure function into a jaxpr and compiles it with XLA
What the code must obey	Almost nothing; uncompiled code is ordinary Python	Python side effects run only while tracing	Pure functions, immutable arrays, explicit random keys, static shapes

In code · python

# PyTorch: eager. Autograd records operations as they run.
import torch
w = torch.randn(3, requires_grad=True)
loss = ((x @ w - y) ** 2).mean()
loss.backward()                    # gradient is now in w.grad

# JAX: functional. grad turns a function into its gradient function.
import jax, jax.numpy as jnp
def loss_fn(w, x, y):
    return jnp.mean((x @ w - y) ** 2)
g = jax.jit(jax.grad(loss_fn))(w, x, y)   # traced once, compiled with XLA

Often confused with

Model serving and MLOps Large language models Running models locally Model hubs JEPA (Joint-Embedding Predictive Architecture)

Where it came from2008–2023

Before frameworks, every new network meant hand-deriving gradients and writing GPU code. Theano, in development from 2008, let researchers write mathematical expressions in Python that it differentiated and compiled for CPUs and GPUs. The whole expression was built first and run later. Google’s TensorFlow, open-sourced under Apache 2.0 in November 2015, took the same approach to a dataflow graph that could run on phones or across hundreds of machines.

PyTorch, incubated at Meta (then Facebook) and released in 2016, went the other way: operations run immediately in ordinary Python, and autograd records the graph as they do. Its 2019 paper argued that this imperative style makes debugging easy and that ease of use and speed need not be traded off. Google Brain’s JAX, open-sourced in 2018, took a third path: keep NumPy-style Python, but trace pure functions and compile them with XLA, treating gradients and compilation as composable transformations.

The styles have since met in the middle. TensorFlow 2 made eager execution the default and kept graphs behind tf.function. PyTorch moved to the Linux Foundation’s new PyTorch Foundation in September 2022, and PyTorch 2.0 added torch.compile in 2023. Keras 3 now runs the same model code on JAX, TensorFlow or PyTorch.

Why it’s built this way

Automatic differentiation instead of hand-written gradients: A model with millions or billions of weights needs a gradient for each one at every step. Recording the operations of the forward pass in order and walking that list backwards gives all of them, for any model the framework can express.
Run eagerly, record as you go: PyTorch chose to make every part of the library ordinary Python so that debuggers, print statements and loops just work. Rebuilding the graph on every iteration is what lets the model’s structure change from one step to the next.
Capture a graph and compile it: Python adds overhead to every small operation, and a compiler can fuse many operations into one kernel. TensorFlow’s guide reports a function running about five times faster as a graph than eagerly, which is why even eager frameworks now offer tf.function, torch.compile and jax.jit.

Timeline

2008
Theano development begins at the University of Montreal
2015
Google open-sources TensorFlow under Apache 2.0 in November
2016
PyTorch is released
2018
JAX is described at SysML and open-sourced by Google Brain
2022
PyTorch moves to the PyTorch Foundation at the Linux Foundation (September 12)
2023
PyTorch 2.0 adds torch.compile (March)

Sources12

Improve this page on GitHub