Deep Learning Frameworks Compared Pytorch Vs Tenso

You’ve just been handed a new project. Maybe it’s a recommendation system, maybe it’s an image classifier. You open your laptop, ready to code, and then it hits you: PyTorch or TensorFlow? And what about this JAX thing everyone keeps mentioning? Your coworker swears by PyTorch’s debugging. A blog post you skimmed says TensorFlow is the only serious option for mobile. Another paper’s code is all in JAX, and it looks like a different language entirely.

It’s confusing. And honestly, it’s not your fault. The framework wars have generated a lot of heat and very little light. Most comparisons are either benchmark shootouts that don’t matter for your project or syntax translations that miss the point entirely. The real differences aren’t about speed—they’re about how you think when you write code.

This article is your decoder ring. We’re not going to crown a winner. Instead, we’re going to build the exact same tiny model in all three frameworks, side-by-side, so you can see the philosophy behind the syntax. By the end, you won’t just know which button to press—you’ll understand why each framework feels the way it does, and you’ll have a clear decision framework for your next project.

Let’s start with the problem that unites them all.

A Single Problem, Three Mindsets: Training a Tiny Model

We need a task that’s big enough to show the moving parts but small enough to hold in your head. Here it is: we’ll train a 2-layer neural network on synthetic regression data. That’s it. No GPUs required, no terabytes of data. Just a model, a loss, and a training loop. The simplicity is the point—it lets us see the forest, not the trees.

We’ll start with PyTorch, because it’s the most familiar to our audience. But as you read the code, pay attention to what the code is forcing you to think about. Every framework is opinionated about three things: how the computation graph gets built, how state (parameters, optimizer momentum) is managed, and how you’re allowed to use Python control flow. Those opinions are what make the frameworks feel like different planets, even when they’re doing the same math.

PyTorch: The Script That Learns

In PyTorch, you define a model as a class. That class owns its parameters—they’re living, mutable objects stored inside the model instance. When you call model(x), you’re running the forward method, and PyTorch is quietly building a computation graph behind the scenes. That graph is built fresh on every single forward pass. This is called “define-by-run,” and it’s the heart of the PyTorch philosophy.

What does define-by-run actually mean? Think of it this way: your training loop is a Python script. It has a for loop, it calls functions, it prints things. PyTorch doesn’t ask you to declare the whole computation upfront. It just follows along, recording every operation as it happens, so it can play the tape backward during backpropagation. The graph is recreated from scratch at every iteration—PyTorch’s official autograd mechanics docs confirm this explicitly.

Let’s see what that looks like in code.

# PyTorch: a 2-layer network on synthetic regression data
# This block is fully self-contained — it generates its own data, defines its
# own model, and runs its own training loop.

import torch
import torch.nn as nn
import torch.optim as optim

# --- Synthetic data ---
# We'll generate 1000 samples with 10 features each. The true relationship is
# linear with a little noise, so a small network can learn it easily.
torch.manual_seed(42)
X = torch.randn(1000, 10)
true_w = torch.randn(10, 1)
y = X @ true_w + 0.1 * torch.randn(1000, 1)

# --- Model definition ---
# The model *owns* its parameters. They're created when we instantiate the
# class, and they live inside `model.parameters()` as mutable tensors.
class TinyNet(nn.Module):
    def __init__(self, in_features=10, hidden=32, out_features=1):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)  # this creates weight & bias
        self.fc2 = nn.Linear(hidden, out_features)

    def forward(self, x):
        # Every call to forward builds a fresh computation graph.
        # PyTorch records each operation so it can compute gradients later.
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = TinyNet()

# --- Optimizer and loss ---
# The optimizer wraps the model's mutable parameters. When we call
# optimizer.step(), it updates those parameters in-place.
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

# --- Training loop ---
# This is pure Python. We're in charge of the loop, the zeroing of gradients,
# the backward pass, and the optimizer step. PyTorch just handles the
# gradient computation when we call loss.backward().
for epoch in range(200):
    # Forward pass: build the graph and compute predictions
    pred = model(X)
    loss = loss_fn(pred, y)

    # Backward pass: compute gradients by playing the graph backward
    optimizer.zero_grad()  # clear old gradients — they accumulate by default!
    loss.backward()        # populate .grad on every parameter
    optimizer.step()       # update parameters in-place using their .grad

    if epoch % 50 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss.item():.4f}")

print(f"Final loss: {loss.item():.4f}")

Let’s interpret what we just saw. The training loop is a for loop—pure Python, nothing magical. optimizer.zero_grad() exists because gradients accumulate by default in PyTorch; if you forget it, your gradients will be the sum of multiple backward passes, which is almost never what you want. loss.backward() triggers the actual gradient computation, walking the graph backward from the loss to every parameter. optimizer.step() updates the parameters in-place.

In plain English: PyTorch feels like writing a script. You’re in charge of the loop, and the framework just handles the gradients. The model’s parameters are stateful objects that get mutated at every step. This is intuitive because it matches how we think about Python objects—they have internal state, and we change that state over time.

Now here’s the interesting part: this intuitive, script-like feeling comes from define-by-run. The graph is ephemeral—it lives for one forward pass and then gets discarded. This is a debugging superpower (we’ll see why later), but it also means PyTorch can’t optimize the graph globally the way a static-graph framework can.

TensorFlow: From Graph-First to Eager-by-Default (and Why That Matters)

To understand TensorFlow, you have to understand its origin story. TensorFlow 1.x was built around a radical idea: you should define your entire computation as a static graph first, and then run it later in a special environment called a Session. The graph was a blueprint, and the session was the factory that executed it.

This had theoretical advantages—the framework could optimize the whole graph before running anything—but it was a debugging nightmare. You couldn’t just print() a tensor’s value, because the tensor was a symbolic placeholder, not a concrete number. You had to run the session, fetch the value, and print it outside the graph. Python control flow like if statements didn’t work naturally either; you had to use special TensorFlow ops like tf.cond.

TensorFlow 2.x threw that out the window. Eager execution became the default. Now you can print() tensors, use Python if statements, and debug normally. The official TensorFlow basics docs frame TF as “an end-to-end platform for machine learning,” and eager mode makes it feel approachable.

But here’s the catch: eager mode is slower. Every operation dispatches to the accelerator individually, which adds overhead. tf.function is the bridge: it traces your Python function, builds a static graph from it, and compiles that graph for speed. The official TensorFlow graph docs show a ~5x speedup for a simple operation (4.1 seconds eager vs. 0.79 seconds graph).

Let’s see the same tiny model in TensorFlow/Keras, first in eager mode, then with @tf.function.

# TensorFlow/Keras: the same 2-layer network, eager mode
# Fully self-contained — generates its own data, defines its own model,
# and runs its own training loop.

import tensorflow as tf
import numpy as np

# --- Synthetic data (same setup as the PyTorch example) ---
np.random.seed(42)
X_np = np.random.randn(1000, 10).astype(np.float32)
true_w_np = np.random.randn(10, 1).astype(np.float32)
y_np = (X_np @ true_w_np + 0.1 * np.random.randn(1000, 1)).astype(np.float32)

# Convert to tf.data.Dataset for efficient batching (even though we use the
# full dataset each step, this is good practice).
dataset = tf.data.Dataset.from_tensor_slices((X_np, y_np)).batch(1000)

# --- Model definition ---
# Keras is the standard high-level API in TensorFlow 2.x. The model *owns*
# its parameters as tf.Variable objects — mutable, just like PyTorch.
model = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1)
])

# --- Optimizer and loss ---
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
loss_fn = tf.keras.losses.MeanSquaredError()

# --- Training loop (eager mode) ---
# In eager mode, operations run immediately. We can print tensors, use
# Python control flow, and debug normally. tf.GradientTape records
# operations for automatic differentiation, similar to PyTorch's autograd.
for epoch in range(200):
    for x_batch, y_batch in dataset:
        with tf.GradientTape() as tape:
            pred = model(x_batch, training=True)
            loss = loss_fn(y_batch, pred)
        # Compute gradients and apply them
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

    if epoch % 50 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss.numpy():.4f}")

print(f"Final loss: {loss.numpy():.4f}")

This looks a lot like the PyTorch version, doesn’t it? A for loop, a forward pass, a gradient computation, an optimizer step. The main difference is tf.GradientTape—you wrap the forward pass in a with block, and the tape records the operations so you can compute gradients afterward. It’s a slightly different API for the same idea.

But now let’s add @tf.function and see what changes.

# TensorFlow/Keras: the same training loop, now with @tf.function
# This block is self-contained — it redefines everything so it can run
# independently and show the timing difference.

import tensorflow as tf
import numpy as np
import time

# --- Synthetic data ---
np.random.seed(42)
X_np = np.random.randn(1000, 10).astype(np.float32)
true_w_np = np.random.randn(10, 1).astype(np.float32)
y_np = (X_np @ true_w_np + 0.1 * np.random.randn(1000, 1)).astype(np.float32)
dataset = tf.data.Dataset.from_tensor_slices((X_np, y_np)).batch(1000)

# --- Model, optimizer, loss (same as before) ---
model = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1)
])
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
loss_fn = tf.keras.losses.MeanSquaredError()

# --- The key change: a @tf.function-decorated training step ---
# When we call this function the FIRST time, TensorFlow *traces* it:
# it runs the Python code, records every TensorFlow operation into a
# static graph, and compiles that graph. Subsequent calls skip the Python
# and run the compiled graph directly on the accelerator.
@tf.function
def train_step(x_batch, y_batch):
    with tf.GradientTape() as tape:
        pred = model(x_batch, training=True)
        loss = loss_fn(y_batch, pred)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# --- Timing comparison ---
# First, a few warm-up steps to trigger tracing and compilation.
for x_batch, y_batch in dataset.take(1):
    _ = train_step(x_batch, y_batch)  # tracing happens here

# Now time the compiled version.
start = time.time()
for _ in range(200):
    for x_batch, y_batch in dataset:
        loss = train_step(x_batch, y_batch)
elapsed = time.time() - start
print(f"200 epochs with @tf.function: {elapsed:.2f}s")
print(f"Final loss: {loss.numpy():.4f}")

What does “tracing” actually mean in plain English? When you call a @tf.function-decorated function for the first time, TensorFlow runs your Python code normally, but it watches every TensorFlow operation you perform. It builds a graph that represents those operations—like a blueprint. Then it compiles that blueprint into efficient accelerator code. On subsequent calls, it skips the Python entirely and just runs the blueprint. That’s where the speedup comes from: you’re no longer paying the overhead of the Python interpreter for every operation.

The 5x speedup the docs mention is real, but for our tiny model, the absolute difference is small because the computation is trivial. The speedup matters when you have complex models and large batch sizes—then the Python overhead becomes a significant fraction of your training time.

Now, here’s the plot twist: Keras 3. As of the official Keras 3 announcement, Keras now runs on JAX and PyTorch too, not just TensorFlow. You can write a model using the Keras API and switch backends with a single line of configuration. The high-level API is converging, even as the low-level execution models diverge. We’ll come back to this.

JAX: Not a Framework — a Function Transformer (And That Changes Everything)

JAX is different. It’s not a deep learning framework in the way PyTorch or TensorFlow are. The JAX quickstart describes it as “array-oriented numerical computation (à la NumPy) with automatic differentiation and JIT compilation.” Notice what’s missing: no mention of neural networks, layers, optimizers, or training loops. JAX gives you NumPy plus three superpowers—grad, jit, and vmap—and says, “build whatever you want on top.”

This is the hardest part of the article, because JAX’s functional-programming model is genuinely unfamiliar if you’re used to PyTorch or TensorFlow. But the core idea is beautiful: you write a pure function that takes parameters and data and returns a loss. Then you transform that function. jax.grad takes your function and returns a new function that computes gradients. jax.jit takes your function and returns a new function that’s compiled for speed. These transformations compose—you can jit(grad(f)) or grad(jit(f)) and it just works. The official “Thinking in JAX” docs emphasize that this composability is the whole point.

The price of admission? Immutability. JAX arrays cannot be changed in-place. You can’t do x[0] = 10. You have to do x.at[0].set(10), which returns a new array. This isn’t a quirk—it’s the load-bearing constraint that makes the transformations work. If arrays could be mutated, the compiler couldn’t reason about your code safely.

Let’s train our tiny model in JAX. We’ll use just JAX plus Optax (a gradient-processing library), no high-level neural-network library like Flax or Haiku. This will expose the functional guts.

# JAX: the same 2-layer network, pure functional style
# Fully self-contained — generates its own data, defines its own
# initialization, loss, and training step, and runs its own loop.

import jax
import jax.numpy as jnp
import optax
import numpy as np

# --- Synthetic data (same setup, but as JAX arrays) ---
np.random.seed(42)
X_np = np.random.randn(1000, 10).astype(np.float32)
true_w_np = np.random.randn(10, 1).astype(np.float32)
y_np = (X_np @ true_w_np + 0.1 * np.random.randn(1000, 1)).astype(np.float32)
X_jax = jnp.array(X_np)
y_jax = jnp.array(y_np)

# --- Parameter initialization ---
# JAX has no nn.Module. We manually initialize parameters as plain dicts
# of JAX arrays. The PRNG key is JAX's explicit random-number generator
# state — no global random seed.
def init_params(rng_key):
    k1, k2 = jax.random.split(rng_key)
    # He initialization-ish: scale by sqrt(2 / fan_in)
    w1 = jax.random.normal(k1, (10, 32)) * jnp.sqrt(2.0 / 10)
    b1 = jnp.zeros(32)
    w2 = jax.random.normal(k2, (32, 1)) * jnp.sqrt(2.0 / 32)
    b2 = jnp.zeros(1)
    return {'w1': w1, 'b1': b1, 'w2': w2, 'b2': b2}

rng = jax.random.PRNGKey(42)
params = init_params(rng)

# --- Forward pass as a pure function ---
# This function takes params and data, returns predictions. No self, no
# internal state. It's a mathematical function: same inputs -> same outputs.
def forward(params, x):
    x = jnp.dot(x, params['w1']) + params['b1']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['w2']) + params['b2']
    return x

# --- Loss as a pure function ---
# Takes params, data, and targets. Returns a scalar loss.
def loss_fn(params, x, y):
    pred = forward(params, x)
    return jnp.mean((pred - y) ** 2)

# --- Optimizer initialization (Optax) ---
# Optax manages optimizer state functionally. We create an optimizer and
# initialize its state from the initial parameters.
optimizer = optax.adam(learning_rate=0.01)
opt_state = optimizer.init(params)

# --- Training step as a pure function ---
# This is the heart of the JAX style. It takes (params, opt_state, data)
# and returns (new_params, new_opt_state, loss). No mutation, no side
# effects. We decorate with @jax.jit to compile it.
@jax.jit
def train_step(params, opt_state, x_batch, y_batch):
    # Compute loss and gradients. jax.value_and_grad returns both the loss
    # value and its gradient with respect to params.
    loss, grads = jax.value_and_grad(loss_fn)(params, x_batch, y_batch)
    # Compute optimizer updates from gradients
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    # Apply updates to get new params (optax.apply_updates returns new params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss

# --- Training loop ---
# Notice: we explicitly pass state in and receive new state out. There is
# no hidden state anywhere. The loop is pure Python, but the step is compiled.
for epoch in range(200):
    params, opt_state, loss = train_step(params, opt_state, X_jax, y_jax)
    if epoch % 50 == 0:
        # .block_until_ready() ensures the async computation is done before
        # we read the value — important for accurate timing, less so here.
        print(f"Epoch {epoch:3d} | Loss: {loss:.4f}")

print(f"Final loss: {loss:.4f}")

Let’s unpack what just happened, because it’s a lot. The training loop has no model object, no optimizer object with internal state. params is a plain dictionary of JAX arrays. opt_state is whatever Optax needs to track momentum and adaptive learning rates. Every call to train_step returns new params and new opt_state—the old ones are not modified. The loop explicitly passes the new state into the next iteration.

This is the hardest part of JAX for newcomers: you manage state like a mathematical function, not like a notebook. In PyTorch, optimizer.step() mutates the parameters in-place, and you don’t have to think about it. In JAX, you thread the state through every function call. It’s more work upfront, but it gives you superpowers. Because train_step is a pure function with no side effects, jax.jit can compile it aggressively. jax.grad can differentiate through it automatically. And jax.vmap can vectorize it across batch dimensions without you writing a single loop.

What about that immutability thing? Let’s see it in action.

# JAX immutability: a quick demonstration
# This block is self-contained.

import jax.numpy as jnp

# Create a JAX array
x = jnp.array([1.0, 2.0, 3.0])
print(f"Original x: {x}")

# This would RAISE AN ERROR:
# x[0] = 10.0  # TypeError: '<class 'jaxlib.xla_extension.ArrayImpl'>' does not support item assignment

# Instead, use .at[index].set(value) to create a NEW array:
x_new = x.at[0].set(10.0)
print(f"New x (after x.at[0].set(10.0)): {x_new}")
print(f"Original x is unchanged:    {x}")

In plain English: JAX arrays are like numbers—you can’t change the number 5 into the number 7; you just compute a new number. This constraint is what lets JAX transform your functions safely. The compiler knows that calling a pure function twice with the same inputs will always give the same outputs, so it can cache computations, fuse operations, and parallelize aggressively.

The Three Philosophies, Side-by-Side: Execution, State, and Control Flow

Now that you’ve seen all three in action, let’s zoom out and compare them along the axes that actually matter for day-to-day work. The table below is a quick reference, but the story is in the prose that follows.

Aspect	PyTorch	TensorFlow 2.x	JAX
Execution model	Define-by-run: graph built fresh every iteration	Eager-by-default; opt-in graph via @tf.function	Pure-functional; explicit JIT via @jax.jit
State management	Mutable objects (nn.Module params, optimizer state)	Mutable objects (tf.Variable, optimizer state)	Immutable arrays; state passed explicitly as function arguments
Control flow	Python if/for always works	Python if/for in eager; must use tf.cond/tf.while_loop in graph mode	Python if/for only outside JIT; inside JIT, use jax.lax.cond/jax.lax.scan
Debugging ease	Easiest: print() and pdb work everywhere	Medium: works in eager, but @tf.function hides values	Hardest: no print() inside JIT; must use jax.debug.print() or return values

Execution Model

PyTorch’s define-by-run means the computation graph is ephemeral—it exists for one forward pass and then vanishes. This is why debugging feels natural: the graph is exactly the sequence of Python operations you just wrote. TensorFlow 2.x starts in eager mode (same feeling) but lets you opt into graph compilation with @tf.function. JAX inverts this: you’re always aware of compilation because you explicitly apply @jax.jit. There’s no hidden graph-building.

State Management

This is where JAX diverges most sharply. In PyTorch and TensorFlow, your model’s parameters are stateful objects—they live inside the model, and the optimizer mutates them in-place. You don’t have to think about state threading. In JAX, parameters are just arrays, and every update produces new arrays. You pass them in, you get new ones out. This is harder at first, but it eliminates entire categories of bugs (like forgetting to zero gradients, or accidentally sharing state between runs).

Control Flow

Here’s the gotcha that bites people. In PyTorch, you can put an if statement anywhere and it just works. In eager TensorFlow, same thing. But inside a @tf.function or a JAX JIT-compiled function, Python control flow is traced—the Python if is evaluated once during tracing, and the branch taken becomes part of the static graph. If you need data-dependent control flow, you have to use framework-specific ops (tf.cond, jax.lax.cond).

Let’s make this concrete with a tiny example.

# Control flow: how a conditional behaves differently in eager vs. compiled mode
# This block demonstrates the tracing gotcha in TensorFlow.

import tensorflow as tf

# --- Eager mode: Python if works as expected ---
def eager_conditional(x):
    # x is a concrete tensor, so the if evaluates normally.
    if tf.reduce_sum(x) > 0:
        return x * 2
    else:
        return x * 3

x = tf.constant([1.0, -2.0])
print(f"Eager result: {eager_conditional(x)}")

# --- Graph mode: Python if is evaluated ONCE during tracing ---
@tf.function
def graph_conditional(x):
    # WARNING: This if is evaluated during TRACING, not during execution.
    # The value of x during tracing determines which branch is baked into
    # the graph FOREVER. If the first traced x has sum > 0, the 'else'
    # branch will NEVER run, even if later inputs have sum <= 0.
    if tf.reduce_sum(x) > 0:
        return x * 2
    else:
        return x * 3

# First call with sum > 0: traces the 'if' branch.
x_pos = tf.constant([1.0, 2.0])
print(f"Graph result (positive): {graph_conditional(x_pos)}")

# Second call with sum <= 0: still uses the traced 'if' branch!
x_neg = tf.constant([-1.0, -2.0])
print(f"Graph result (negative): {graph_conditional(x_neg)} — OOPS, should be x*3!")

# The correct way: use tf.cond for data-dependent control flow in graph mode.
@tf.function
def correct_graph_conditional(x):
    return tf.cond(
        tf.reduce_sum(x) > 0,
        lambda: x * 2,
        lambda: x * 3
    )

print(f"Correct graph (positive): {correct_graph_conditional(x_pos)}")
print(f"Correct graph (negative): {correct_graph_conditional(x_neg)}")

This is the kind of silent bug that drives people crazy. The code looks like normal Python, but @tf.function changes the rules. JAX has the same issue: inside a JIT-compiled function, Python control flow is traced, and you need jax.lax.cond for data-dependent branches.

The Keras 3 Convergence

Before we move on, let’s acknowledge the elephant in the room: Keras 3 now runs on JAX and PyTorch too. At the high level—the level where you stack Dense layers and call model.fit()—all three frameworks now share a common API. You can write a model in Keras and run it on any backend. The differences we’ve been discussing live at the lower level: custom training loops, research code, performance-critical components. If you never write a custom train_step, you might not feel the differences at all. But when you do, they’re profound.

Debugging and the “It Works, But Why?” Problem

Your model isn’t training. The loss is flatlining. Is it vanishing gradients? A bug in the data pipeline? A shape mismatch that’s silently broadcasting? How you answer that question depends dramatically on which framework you’re using.

PyTorch: The Debugging Superpower

In PyTorch, you can stick a print() anywhere. Inside the forward pass, inside the training loop, even inside a gradient hook. pdb.set_trace() drops you into a debugger where you can inspect every tensor. The dynamic graph means what you see is what you get—the tensor you’re looking at is the actual value flowing through the network right now.

# PyTorch debugging: printing gradients mid-training
# Fully self-contained.

import torch
import torch.nn as nn
import torch.optim as optim

# Quick setup (same as before, abbreviated)
torch.manual_seed(42)
X = torch.randn(100, 10)
y = torch.randn(100, 1)

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

# Training step with debugging
pred = model(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()

# We can print gradients directly — they're just tensors.
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad mean = {param.grad.mean().item():.6f}, "
              f"grad std = {param.grad.std().item():.6f}")

# We can also use pdb to drop into an interactive debugger:
# import pdb; pdb.set_trace()  # uncomment to try

TensorFlow: Eager Is Easy, Graph Is Hard

In eager mode, TensorFlow debugging feels like PyTorch. But once you add @tf.function, print() only runs during tracing—not during execution. If you want to see a tensor’s value at runtime, you need tf.print(). This is a common source of silent bugs: you add a print() for debugging, it runs once during tracing and never again, and you think your code is fine.

# TensorFlow debugging: the @tf.function print gotcha
# Fully self-contained.

import tensorflow as tf

@tf.function
def buggy_debug(x):
    # This Python print runs ONLY during tracing, not every call.
    print(f"Tracing! x shape: {x.shape}")
    # To print at runtime, use tf.print:
    tf.print("Runtime x sum:", tf.reduce_sum(x))
    return x * 2

x = tf.constant([1.0, 2.0, 3.0])
print("First call:")
_ = buggy_debug(x)
print("Second call (notice no 'Tracing!' message):")
_ = buggy_debug(x)

JAX: No Python Inside the Compiled Function

JAX is the hardest to debug because JIT-compiled functions run entirely on the accelerator with no Python interpreter. You can’t print() inside them at all. You have to use jax.debug.print() (which works at runtime) or extract values by returning them from the compiled function. The functional style also means you can’t just “peek” at a variable—you have to return it.

# JAX debugging: printing inside a JIT-compiled function
# Fully self-contained.

import jax
import jax.numpy as jnp

@jax.jit
def debug_step(x):
    # This Python print would run during tracing, not execution.
    # print(x)  # Uncomment to see: it prints a tracer, not a value.

    # The correct way: use jax.debug.print for runtime values.
    jax.debug.print("Runtime x sum: {s}", s=jnp.sum(x))
    return x * 2

x = jnp.array([1.0, 2.0, 3.0])
print("Calling debug_step:")
result = debug_step(x)
print(f"Result: {result}")

This is the hardest part of JAX for newcomers: the debugging workflow is fundamentally different, and the error messages from XLA (the compiler) can be cryptic. But here’s the silver lining: once your code works, it’s guaranteed to be free of certain classes of side-effect bugs. The functional constraints that make debugging harder also make the code more predictable.

Ecosystem and Deployment: Where the Rubber Meets the Road

Writing a model is one thing. Shipping it is another. And this is where the frameworks diverge sharply—not in what they can do, but in how smooth the path is.

TensorFlow has the most mature deployment story. TensorFlow Serving lets you deploy models as production-grade gRPC or REST services. TensorFlow Lite compiles models for mobile and embedded devices. TensorFlow.js runs models in the browser. If your end goal is a model in a phone app or a web page, TensorFlow has the smoothest, most-documented path. The official docs frame TF as “an end-to-end platform,” and deployment is where that platform pitch still has teeth.

PyTorch has caught up significantly. TorchServe provides model serving. PyTorch Mobile brings models to iOS and Android. ONNX export lets you move models between frameworks. But the ecosystem is more fragmented—you’ll often find yourself stitching tools together rather than following a single paved path.

JAX has no built-in serving solution. You typically export to TensorFlow SavedModel or ONNX, or run the model in-process behind a lightweight server like FastAPI. This is fine for research and internal services, but it’s a real gap for production teams that need battle-tested serving infrastructure.

Deployment Target	TensorFlow	PyTorch	JAX
Server (REST/gRPC)	TensorFlow Serving (mature)	TorchServe (maturing)	Export to TF SavedModel or use FastAPI in-process
Mobile (iOS/Android)	TensorFlow Lite (mature)	PyTorch Mobile (maturing)	Export to TFLite or ONNX
Browser (JavaScript)	TensorFlow.js (mature)	ONNX.js or TF.js via export	Export to TF.js via SavedModel
Edge / Embedded	TensorFlow Lite Micro	ExecuTorch (new)	Limited

Here’s the practical bridge: Keras 3. You can prototype in PyTorch, train at scale with JAX’s JIT compilation, and deploy with TensorFlow Serving—all from the same Keras model definition. The high-level API convergence means you’re not locked into one framework’s deployment story anymore.

Which One Should You Actually Use? A Decision Framework, Not a Winner

You’ve seen the code. You’ve seen the philosophies. Now you want an answer. But the right answer depends on your context, not on a universal “best” framework. Here’s a decision framework based on who you are and what you’re building.

If you’re learning deep learning for the first time: Start with PyTorch. The debugging experience is the most forgiving—you can print() anything, use pdb normally, and the dynamic graph means what you see is what you get. The community is the largest, so you’ll find more tutorials, more answered StackOverflow questions, and more colleagues who can help. PyTorch’s define-by-run philosophy matches how you already think about Python code.

If you’re on a product team that needs to deploy to mobile or web: TensorFlow (with Keras) still has the smoothest path. TensorFlow Lite and TensorFlow.js are mature, well-documented, and battle-tested. The gap is narrowing—PyTorch Mobile and ExecuTorch are improving fast—but as of today, TF’s deployment story is the most complete.

If you’re doing research that pushes the boundaries of parallelism or compilation: JAX’s functional transformations give you superpowers that are hard to replicate in PyTorch or TensorFlow. If you’re training large-scale transformers, doing meta-learning, or writing custom parallel training loops, JAX’s pmap and vmap let you express parallelism in a way that feels natural once you’ve internalized the functional style. The learning curve is real, but the payoff is real too.

If you want to keep your options open: Use Keras 3. Write your model once, and switch backends as your needs evolve. Prototype in PyTorch for easy debugging, switch to JAX for fast training, deploy with TensorFlow Serving. Keras 3 is the practical manifestation of the convergence we’ve been talking about.

The bottom line: The frameworks are converging at the high level (Keras) and diverging at the low level (execution models). The choice matters most when you’re writing custom training loops or performance-critical components—and less when you’re stacking standard layers. Pick the framework that matches your brain and your deployment target, and know that switching is getting easier every year.

Recap: What We Learned

• PyTorch feels like writing a Python script. It’s define-by-run, so the computation graph is built fresh every iteration. Debugging is straightforward—print() and pdb work everywhere. The model owns its parameters as mutable objects.
• TensorFlow 2.x is eager-first, with a speed boost when you want it. @tf.function traces your Python code into a static graph, giving you a ~5x speedup at the cost of debugging complexity. Keras 3 now runs on JAX and PyTorch too, making the high-level API a convergence point.
• JAX is a function transformer, not a framework. You write pure functions and transform them with grad, jit, and vmap. Immutability and explicit state-passing are the price of admission, but they unlock powerful compilation and parallelization.
• The frameworks are converging at the high-level API and diverging at the low-level execution model. The choice matters most when you write custom training code.
• Debugging is the hidden differentiator. PyTorch is the easiest to debug, JAX is the hardest, and TensorFlow sits in the middle with @tf.function gotchas.

Check Your Understanding

Remember: What are the three superpowers JAX adds on top of NumPy? Name them and give a one-sentence plain-English description of each.

Understand: In your own words, explain what “define-by-run” means in PyTorch and why it makes debugging easier. Contrast it with TensorFlow 1.x’s static-graph approach.

Apply: You’re debugging a TensorFlow training loop where the loss isn’t decreasing. You add a print(loss) inside a @tf.function-decorated training step. What will happen, and why might it mislead you? What should you use instead?

Analyze: Compare how PyTorch and JAX handle parameter updates during training. Why does JAX’s functional approach eliminate the need for optimizer.zero_grad()?

Evaluate: A startup is building a mobile app that runs an image classifier on-device. They have equal expertise in all three frameworks. Which framework would you recommend for deployment, and why? What tradeoffs are they making?

Create: Write a JAX training step function that includes L2 regularization. The function should take params, opt_state, x_batch, y_batch, and a l2_lambda scalar, and return the updated state and loss. Explain why the regularization term must be added inside the loss function rather than applied as a separate mutation.