Understanding JAX means understanding how it differs philosophically from the other major frameworks. Let's be precise about these differences — and honest about when another tool might be the better choice.
PyTorch: Eager + Autograd Tape
PyTorch uses eager execution with an autograd tape. When you run operations, they execute immediately and PyTorch records them on a computational tape. When you call .backward(), it replays the tape in reverse to compute gradients. The programming model is imperative and object-oriented: you define models as classes inheriting from nn.Module.
# PyTorch style (for comparison)
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(2, 1)
def forward(self, x):
return self.linear(x)
model = Model()
x = torch.tensor([[1.0, 2.0]])
y = model(x) # Eager execution, tape is recording
loss = y.sum()
loss.backward() # Replay tape to get gradients
print(model.linear.weight.grad)
TensorFlow: Computational Graphs
TensorFlow (v2 with eager mode, or v1's graph mode) historically used define-then-run computational graphs. TF2 made eager execution the default but retained tf.function for graph compilation. State is managed through tf.Variable objects. The programming model mixes functional and object-oriented styles.
JAX: Pure Function Transformations
JAX takes a functional approach. There are no model classes, no autograd tape, no built-in training loops. You write pure functions, and JAX provides transformations to differentiate, compile, and parallelize them.
# JAX style
import jax
import jax.numpy as jnp
def predict(params, x):
w, b = params
return jnp.dot(x, w) + b
def loss_fn(params, x, y):
pred = predict(params, x)
return jnp.mean((pred - y) ** 2)
# Parameters are just arrays in a tuple — no special Variable type
params = (jnp.array([0.5, 0.3]), jnp.array(0.1))
x = jnp.array([[1.0, 2.0]])
y = jnp.array([1.5])
# Gradient is a function, not a method on a loss object
grads = jax.grad(loss_fn)(params, x, y)
print(grads) # Tuple of gradient arrays matching params structure
💡 Why This Matters
The functional style isn't just aesthetic preference — it has real consequences. Because JAX functions are pure, the compiler can aggressively optimize them. Because parameters are explicit function arguments (not hidden object state), it's trivial to do things like compute gradients with respect to different parameter subsets, or checkpoint/serialize state.
When to Choose Each Framework
| Choose JAX when... | Choose PyTorch when... |
|---|---|
| You need high-performance research code | You need mature deployment tools (TorchServe, ONNX) |
| You're doing scientific computing or physics | Your team already knows PyTorch |
| You want TPU support as a first-class citizen | You need the largest ecosystem of pretrained models |
| You need composable transforms (vmap, etc.) | You prefer object-oriented model building |
| You're at a research lab that uses JAX | You need extensive community tutorials and SO answers |
Honest take: PyTorch has a much larger ecosystem, more tutorials, and more pretrained model hubs. If you're building a standard CNN or transformer for a production app, PyTorch or its ecosystem (HuggingFace, etc.) is often the path of least resistance. JAX shines when you need something PyTorch can't easily give you — custom differentiation patterns, seamless TPU scaling, or scientific computing workloads where you're not using standard neural network layers.