PyTorch, TensorFlow and JAX Compared

Why the framework you choose still matters

PyTorch, TensorFlow and JAX all let you build neural networks, compute gradients automatically, and run computation on GPUs or TPUs. At that level they do the same job, and a model trained in one can often be reproduced, with effort, in either of the others. But the three frameworks disagree about how code should run, what it feels like to hold a tensor in your hands, and how a training loop should be structured. Those disagreements surface in nearly every line you write, in how you debug a shape mismatch, and in how you eventually get a trained model into production, so the choice is not just a matter of taste.

TensorFlow was the first of the three, built by Google around the idea of first describing a computation as a graph and then running that graph with real data. PyTorch followed, built by Meta (then Facebook), around the opposite idea: code runs immediately, the way ordinary Python does, and the graph is built quietly behind the scenes as a side effect. JAX came later, also associated with Google, not really as a deep learning framework but as a fast, NumPy-flavoured library with automatic differentiation and compilation added as general-purpose tools rather than deep-learning-specific features. Over time the three have borrowed from each other. TensorFlow added eager execution so it feels more like PyTorch day to day. PyTorch added a compiler so it can get graph-like speed without giving up its eager style. JAX libraries added the layers, optimizers and training utilities that turn a general array library into something you can train a model with. The capability gap between them has narrowed considerably. What has not narrowed is the underlying mental model each one expects you to adopt, and that is where this comparison starts.

Three different ideas about how code should run

PyTorch: eager by default, graphs are optional

In PyTorch, every line of code runs the moment Python reaches it, the same way plain NumPy does. Write y = x * 2 and that multiplication happens immediately; y is a real tensor holding real numbers that you can print, index into, or pass to an ordinary if statement. Underneath, PyTorch records each operation performed on a tensor that needs a gradient, quietly building a graph as a side effect of running your code. When you later call .backward() on a result, PyTorch walks that recorded graph backwards and computes gradients with respect to every tensor that fed into it. This style is often called define-by-run: the graph is defined by the act of running the code, and a fresh graph is built on every single call. The benefit is that debugging feels completely normal, since you can drop in a print statement or a debugger breakpoint anywhere and see real values, and models whose structure changes between inputs, such as recursive structures or networks with data-dependent branching, are straightforward to write because the graph simply reflects whatever path the code actually took.

TensorFlow: graphs first, eager later, tf.function now

The original version of TensorFlow worked the other way around. You first wrote a description of the entire computation as a static graph, with placeholders standing in for data that did not exist yet, and only afterwards fed real numbers through that graph inside a session. This define-and-run approach made certain optimizations and deployment scenarios easier to support, but it made everyday development harder, since errors could surface in a place far removed from the line of code that actually caused them, and simple tasks like printing an intermediate value required special graph operations rather than an ordinary print statement. TensorFlow 2 changed the default to eager execution, so most day-to-day code now looks and behaves like PyTorch, with operations running immediately and values visible as you go. The graph idea did not disappear, it moved behind an explicit opt-in. Decorating a Python function with the tf.function decorator tells TensorFlow to trace that function once, convert it into a graph, and reuse the compiled version on subsequent calls, which is how you recover speed once a piece of code has stabilised and no longer needs to be stepped through line by line.

JAX: NumPy plus function transformations

JAX does not set out to be a deep learning framework at its core, it is a library for writing ordinary NumPy-style array code and then applying transformations to that code. jax.grad takes a function and returns a new function that computes its gradient. jax.jit takes a function and returns a compiled version of it. jax.vmap takes a function written to handle one example and returns a version that runs over a whole batch, without you writing the batching logic by hand. jax.pmap spreads a computation across multiple accelerators. These transformations are composable, so you can jit a function produced by grad, or vmap a function produced by jit. JAX arrays are immutable, there is no way to change one in place, every operation produces a new array, and the functions you pass into these transformations are expected to be pure, meaning they have no hidden side effects and their output depends only on their inputs, not on any external state. Because JAX itself only provides arrays and transformations, building an actual neural network means reaching for a separate library on top, typically Flax or Haiku for defining layers and model state, and Optax for optimizers such as SGD and Adam.

Tensors, devices and the little things that differ day to day

Before getting into gradients, it helps to see how each framework handles the basic act of creating an array and deciding where it lives, CPU, GPU or TPU, because this is where a lot of early confusion happens. PyTorch calls its array type torch.Tensor. TensorFlow uses tf.Tensor for ordinary values and tf.Variable for values meant to change over time, such as model weights. JAX arrays are created through jax.numpy and behave like NumPy arrays, supporting the same indexing, broadcasting and slicing rules you already know if you have used NumPy.

import torch
a = torch.tensor([1.0, 2.0, 3.0])
a = a.to("cuda")  # moves the tensor to the GPU; raises an error if no GPU is visible

import tensorflow as tf
b = tf.constant([1.0, 2.0, 3.0])
with tf.device("/GPU:0"):
    c = b * 2  # runs this operation on the GPU

import jax.numpy as jnp
d = jnp.array([1.0, 2.0, 3.0])
# JAX places this on the default accelerator automatically if one is visible to it

PyTorch requires you to explicitly move a tensor to a device with .to(device) or the shortcut .cuda(), and an operation between two tensors that live on different devices raises an error, which is explicit but means you always know exactly what runs where. TensorFlow behaves similarly with its tf.device context manager, though it will also place operations on an available GPU automatically if you never specify a device. JAX detects whatever accelerator is visible on the machine, GPU or TPU, and places newly created arrays there by default. This is convenient, but it also means a JAX script can quietly already be running on a GPU with no code changes at all, while an equivalent PyTorch script sitting next to it is still running on the CPU because nobody called .cuda(), which is worth knowing before concluding that one framework is faster than another purely by watching a script run.

The same tiny problem, three ways

The clearest way to see these different mental models is to solve the smallest possible optimization problem in each framework: find the value of x that minimizes f(x) = (x - 3) squared, using five steps of plain gradient descent starting from x = 0. The true minimum sits at x = 3, where the slope of f is zero. The examples below assume PyTorch 2.x, TensorFlow 2.x, and a recent release of jax and jaxlib; the core API shown here has been stable across recent versions of all three.

import torch

x = torch.tensor(0.0, requires_grad=True)
lr = 0.1

for step in range(5):
    y = (x - 3) ** 2
    y.backward()
    with torch.no_grad():
        x -= lr * x.grad
        x.grad.zero_()
    print(step, x.item())

x is created with requires_grad=True so PyTorch tracks operations performed on it. Calling y.backward() computes dy/dx and stores the result in x.grad. The update itself has to happen inside a torch.no_grad() block, because otherwise the subtraction would itself be recorded as part of the graph. The gradient also has to be reset to zero after each step with x.grad.zero_(), because PyTorch accumulates gradients into .grad by default rather than overwriting them, which is one of the most common sources of subtle bugs for people new to the framework.

import tensorflow as tf

x = tf.Variable(0.0)
lr = 0.1

for step in range(5):
    with tf.GradientTape() as tape:
        y = (x - 3) ** 2
    grad = tape.gradient(y, x)
    x.assign_sub(lr * grad)
    print(step, x.numpy())

tf.Variable marks x as something whose value can change over time and whose use should be tracked. The tf.GradientTape() block records every operation performed inside it, so the tape can later differentiate y with respect to x. Outside the tape, x.assign_sub updates the variable in place. No manual resetting is needed here, because the tape is a fresh, temporary recording created on each loop iteration rather than a persistent graph attached to x.

import jax
import jax.numpy as jnp

def f(x):
    return (x - 3) ** 2

grad_f = jax.grad(f)

x = 0.0
lr = 0.1

for step in range(5):
    g = grad_f(x)
    x = x - lr * g
    print(step, x)

f is an ordinary Python function with no special tensor type involved. jax.grad(f) does not compute a gradient itself, it returns a new function, grad_f, that computes the gradient of f at whatever point you call it with. There is no object called x that carries a gradient attribute, x is just a plain Python float that gets reassigned on each loop iteration. Nothing is tracked or recorded on x; the tracking happens inside grad_f at the moment it is called.

All three snippets converge towards x = 3, but they encode three different mental models for the same five lines of math. PyTorch attaches a graph to the tensor and accumulates gradients on it until you clear them. TensorFlow records a temporary tape around a block of code and discards it once that block ends. JAX treats differentiation as a transformation of a function, with no state attached to the value being optimized at all. Once this distinction clicks, most of the remaining differences between the frameworks turn out to be details built on top of it.

Automatic differentiation under the hood

All three frameworks use reverse-mode automatic differentiation, which is the same underlying technique regardless of the API wrapped around it. The idea is to break a computation into a sequence of elementary operations, such as additions, multiplications and standard functions, and track how each one depends on the values that fed into it. A forward pass computes the actual output using real numbers. A backward pass then walks the sequence of operations in reverse, applying the chain rule at each step to propagate a derivative from the output back towards the original inputs. This is efficient precisely when a function has one scalar output and many inputs, which is exactly the shape of a loss function computed from millions of model parameters, since one backward pass produces the gradient with respect to every parameter at once.

Where the frameworks differ is in when and how that sequence of operations gets recorded. PyTorch builds a dynamic graph attached to tensors on every forward pass and discards it after backward() unless you explicitly ask it to keep the graph around. TensorFlow's GradientTape similarly builds a temporary graph scoped to whatever code runs inside the tape block, and by default only tracks tf.Variable objects, so a plain tf.constant needs an explicit tape.watch() call if you want a gradient with respect to it. JAX traces a function the moment you call jax.grad or jax.jit on it, and by default jax.grad only works on functions that return a single scalar value; to get a full Jacobian for a function with multiple outputs you reach for jax.jacfwd or jax.jacrev instead.

• PyTorch accumulates gradients in .grad by default, so training loops call optimizer.zero_grad() or tensor.grad.zero_() between steps, and forgetting this is a common source of erratic training.
• TensorFlow's GradientTape only tracks tf.Variable by default; differentiating with respect to a plain tensor requires an explicit tape.watch() call inside the tape block.
• JAX's grad expects the function to return a scalar; applying it directly to a function with a vector or matrix output raises an error rather than returning something unexpected.
• All three can compute higher-order derivatives (a gradient of a gradient), but the syntax for it differs, and it is worth checking the current documentation for whichever framework you use rather than guessing.

Making code fast: compilation and XLA

Running Python eagerly, one operation at a time, is flexible but comes with overhead, since Python itself has to dispatch each small operation individually and the underlying engine cannot see ahead to optimize across several operations at once. Compiling a function into a fixed graph lets the engine fuse operations together, plan memory use ahead of time, and in many cases run noticeably faster on an accelerator. All three frameworks offer a route to this kind of compilation, and two of them share the same compiler underneath.

torch.compile

PyTorch 2.0 introduced torch.compile, which lets you keep writing ordinary eager-style code and simply wrap a model or function, for example model = torch.compile(model), to get many of the speed benefits of a graph-based approach without restructuring your code into a different style.

tf.function and XLA

TensorFlow's tf.function traces Python code into a graph the first time it runs. Passing jit_compile=True to tf.function hands that traced graph to XLA, a compiler that fuses operations and generates code tailored to the specific device, which is the same compiler that sits underneath JAX.

jax.jit

JAX compiles through XLA by design rather than as an optional extra. jax.jit(f) traces f once for each distinct combination of input shapes and data types and caches the compiled version, reusing it on later calls with matching shapes. Because JAX expects pure functions, ordinary Python side effects inside a jitted function behave in surprising ways. A print statement inside a function wrapped with jax.jit, for instance, can show a placeholder description of a traced array rather than an actual number, because tracing sometimes runs with abstract stand-ins for values instead of the concrete numbers themselves. This catches a lot of newcomers the first time they try to debug a jitted function with print statements.

A practical warning applies to all three: the very first call to a compiled or jitted function pays a one-off tracing and compilation cost that later calls do not. Timing only the first call and concluding that a framework is slow, or timing it against an eager call in another framework, is a common and avoidable mistake. Fair comparisons warm up the function with a call or two first, then time several repeated calls.

Ecosystem: layers, training loops and pretrained models

PyTorch's ecosystem includes torchvision, torchaudio and torchtext for common data loading and pretrained models in their respective domains. A large share of newly published research code is released in PyTorch, which matters if your work involves reproducing results from recent papers. The Hugging Face Transformers library, widely used for pretrained language and vision models, was originally built PyTorch-first, though it now supports other backends as well. Higher-level libraries such as PyTorch Lightning and the Accelerate library handle a lot of the boilerplate around training loops, logging and multi-device training so you do not have to write it from scratch.

TensorFlow bundles Keras as tf.keras, a high-level API that makes it quick to assemble standard architectures from predefined layers, and this remains one of the most approachable entry points into deep learning for beginners. A separately maintained release called Keras 3 can run on top of TensorFlow, PyTorch or JAX as a backend, which means some of the boundaries described in this article are softer than they used to be, since the same Keras code can in principle execute on any of the three engines. TensorFlow Extended, or TFX, provides tooling for building production data and training pipelines end to end.

JAX does not bundle a single high-level API of its own. Neural network layers typically come from Flax or Haiku, optimizers come from Optax, and checkpointing and serialization are often handled with a separate tool such as Orbax. This modularity means more setup for a first project, since you are choosing and combining several libraries rather than importing one, but it also means more flexibility, and it is the toolkit of choice for a significant amount of large-scale research associated with Google and DeepMind.

Getting models into the real world: deployment

PyTorch offers TorchScript, which traces or scripts a model into a form that can run without a Python interpreter, along with support for exporting models to ONNX, an open format that other runtimes can load, and PyTorch Mobile for running models on phones.

TensorFlow has historically had the strongest deployment story of the three, built with production in mind from early on. The SavedModel format, TensorFlow Serving for serving models from a server, TensorFlow Lite for mobile and embedded devices, and TensorFlow.js for running models directly in a browser together cover a wide range of deployment targets with first-party tooling.

JAX has fewer dedicated deployment tools of its own. A common pattern is to train a model in JAX and then export its parameters, or convert the computation through ONNX or another runtime, for serving outside of a JAX process, or to serve it directly through a lightweight JAX or XLA runtime if the surrounding infrastructure already supports that.

Common mistakes when comparing or switching

• Forgetting to zero gradients in PyTorch: because .grad accumulates by default, skipping optimizer.zero_grad() between steps silently adds each new gradient on top of the last one, producing training that looks unstable for no obvious reason.
• Treating JAX arrays as mutable: writing array[0] = 1 on a JAX array raises an error, since JAX arrays are immutable; the equivalent is array.at[0].set(1), which returns a new array rather than changing the old one in place.
• Using data-dependent Python control flow inside a jax.jit function: an if statement that branches on the value of a traced array either fails outright or silently follows whichever branch was taken during the one-time trace, regardless of the actual input later.
• Benchmarking the first call to a compiled or jitted function: the first call includes a one-off tracing and compilation cost, so timing only that call and comparing it to an eager call elsewhere overstates how slow the compiled approach is.
• Expecting identical numerical results across frameworks: differences in default data types, the order operations are carried out in, and reduction algorithms mean results can differ slightly, often in a decimal place far enough along that it does not matter in practice but is still worth expecting.
• Treating the three as permanently separate categories: tools like Keras 3, which runs on top of TensorFlow, PyTorch or JAX, and export formats like ONNX mean the practical boundaries between the frameworks are softer than a strict comparison suggests.

Which framework should you learn first

The honest answer depends on what you are trying to do next, more than on which framework is objectively best, since all three can train the same kinds of models to the same accuracy given enough effort.

• Reading and reproducing recent research papers: PyTorch, since most published research code is released in it.
• Working inside an existing TensorFlow codebase, or deploying to mobile or the browser: TensorFlow, since its deployment tooling is the most mature of the three.
• Large-scale training research, custom optimizers, or working somewhere that already standardizes on it: JAX, where the extra setup pays off in control over compilation and parallelism.
• No specific constraint and this is your first framework: PyTorch's eager style and large ecosystem make it a reasonable default to start with.

Whichever one you start with, the underlying ideas, tensors, broadcasting, automatic differentiation, gradient descent, and graphs built either implicitly or explicitly, transfer directly to the other two. Learning a second framework after the first is mostly a matter of learning new names for concepts you already understand, which is a far smaller task than learning them for the first time.

Summary and what comes next

PyTorch, TensorFlow and JAX all solve the same core problem: representing arrays, computing gradients automatically, and running computation efficiently on accelerators. They differ in the assumptions they make about how code should run. PyTorch runs eagerly and builds a graph as a side effect of execution, discarding and rebuilding it on every call. TensorFlow now also runs eagerly by default but can trace a function into a reusable graph with tf.function when speed matters. JAX treats both differentiation and compilation as explicit transformations applied to pure, NumPy-style functions, with no hidden state attached to the values themselves. Once that distinction is clear, most of what looks like a difference in capability turns out to be a difference in API and philosophy sitting on top of the same mathematics.

A good next step is to implement the five-line gradient descent example from this article in whichever framework is newest to you, then extend it to a small neural network on a toy dataset, training for a few dozen steps. Pay attention to where each framework's mental model causes a different kind of bug, a forgotten zero_grad in PyTorch, an unwatched tensor in TensorFlow, or an attempted in-place update in JAX. That hands-on friction teaches the differences between these three frameworks more effectively than any side-by-side comparison can.