ML Frameworks Compared: TensorFlow, PyTorch & JAX

Three philosophies, one job

TensorFlow, PyTorch, and JAX all do the same fundamental job: let you define a numerical computation involving tensors, compute gradients of that computation automatically, and run it efficiently on GPUs or TPUs. What separates them is not capability so much as philosophy — how the computation is expressed, when it is compiled, and who the framework imagines is using it.

Understanding those philosophical differences explains almost everything else: why PyTorch dominates research, why TensorFlow still runs in production pipelines you've never heard of, and why JAX quietly became the backbone of most large-scale model training at Google and several frontier labs.

Define-and-run vs. eager execution

Early TensorFlow (1.x) used a define-and-run (also called "graph mode" or "static graph") model. You first built a symbolic computation graph — a blueprint describing operations without executing them — and then ran data through that fixed graph in a separate session. This let TensorFlow aggressively optimize the whole graph before execution, but debugging was painful: a stack trace pointed at graph-construction code, not the line that actually produced a bad value.

PyTorch popularized the opposite approach: eager execution. Operations run immediately, line by line, exactly like normal Python. print(tensor) works anywhere. You can set a breakpoint inside a training loop and inspect real values. This made PyTorch dramatically easier to debug and iterate with, which is precisely why it won over the research community so quickly after its 2017 release.

TensorFlow 2.x responded by making eager execution the default too, with tf.function available to trace and compile a function into a graph when you want the speed. In other words, the two frameworks converged toward the same middle ground from opposite starting points.

JAX: functional purity as a design constraint

JAX takes a different stance entirely. Instead of an object-oriented API with mutable model objects (model.forward(), optimizer.step()), JAX models functions as pure functions of their inputs: no hidden state, no side effects, same output every time for the same input. This isn't a stylistic preference — it's what makes JAX's core transformations possible.

JAX is built around a small set of function transformations: grad (automatic differentiation), jit (just-in-time compilation via XLA), vmap (automatic vectorization — turning a function written for one example into one that runs on a batch, without manual batching logic), and pmap/shard_map (parallelization across devices). Because functions are pure, these transformations can be composed freely: jit(vmap(grad(loss_fn))) is a perfectly normal thing to write.

The cost is a steeper learning curve — no in-place mutation, explicit random number key management instead of global RNG state, and a functional style that feels unfamiliar if you came from NumPy-and-objects habits. The payoff is a cleaner path to large-scale, highly parallel training.

TensorFlow vs. PyTorch vs. JAX at a glance

How the landscape consolidated

Five years ago this was a genuine three (or four, counting Theano's legacy and MXNet) horse race. By the mid-2020s it consolidated sharply: PyTorch became the default for research and for most new model training, including the large majority of models released on Hugging Face. Hugging Face's own transformers library, originally multi-backend, now treats PyTorch as the primary target.

TensorFlow didn't disappear — it receded into production infrastructure. Companies with mature TF-based serving pipelines, mobile deployment via TFLite, or TensorFlow.js in the browser kept it, because rewriting a stable production system has no payoff. New projects, though, rarely start there.

JAX carved out a durable niche at the frontier of scale. Google's largest models, several of Anthropic's and DeepMind's research lines, and performance-sensitive research labs use JAX specifically because its pmap/shard_map model and tight XLA integration make multi-thousand-chip training setups more tractable to express correctly. It never became a mainstream choice for day-to-day applied work, but it didn't need to — its audience was always systems at the extreme end of scale.

Which one is right today

A pragmatic default for 2026: start in PyTorch unless you have a specific reason not to. It has the largest ecosystem, the most pretrained models, the gentlest debugging experience, and torch.compile now closes most of the performance gap that used to justify reaching for something else.

Reach for JAX when you are training at a scale where XLA's compiler optimizations and JAX's parallelism primitives materially change what's feasible — typically very large models across many accelerators — or when your research requires custom, composable differentiation (e.g., meta-learning, differentiating through an optimizer).

Stay on TensorFlow only if you're maintaining an existing TF production system, need TFLite for mobile/edge deployment, or your deployment story is already built around the TF ecosystem.

What's next

Framework choice determines how you express computation; it doesn't determine what hardware that computation runs on or how fast it goes. The next lesson, ML Hardware Acceleration: GPUs, TPUs & Accelerators, looks at the actual chips underneath every one of these frameworks and why matrix multiplication shapes the entire hardware landscape.