JAX for NumPy Users: What's Actually Different (and Why It's Fast)

You have a 1000×1000 matrix multiplication. In NumPy, it takes half a second. You’ve heard JAX is faster — Google’s new library, used by DeepMind, built for speed. You swap np for jnp and run it again. And… it’s slower. Not by a little — by a lot. What’s going on?

The answer is both frustrating and fascinating. A naive swap from NumPy to JAX can actually make your code slower because of what’s called dispatch overhead — the work JAX does behind the scenes before it even starts computing. But with one tiny change (wrapping your function in jax.jit), that same 1000×1000 matmul drops from ~0.543 seconds to ~0.035 seconds. That’s a 15x speedup.

The catch? That one-liner only works if you understand what JAX is actually doing under the hood. And that’s the whole point of this article. By the end, you’ll understand why JAX is fast, when it isn’t, and how to use its real superpowers — automatic differentiation, auto-vectorization, and multi-device parallelism — safely and effectively.

Let’s start with the benchmark that proves the mystery exists.

# Self-contained benchmark: NumPy vs JAX (uncompiled) vs JAX (JIT)
import numpy as np
import jax.numpy as jnp
from jax import jit
import time

# Create large random matrices
n = 1000
np.random.seed(42)
a_np = np.random.randn(n, n)
b_np = np.random.randn(n, n)

a_jax = jnp.array(a_np)
b_jax = jnp.array(b_np)

# NumPy timing
start = time.time()
_ = a_np @ b_np
numpy_time = time.time() - start
print(f"NumPy: {numpy_time:.3f}s")

# JAX uncompiled timing (must call .block_until_ready() to get real time)
start = time.time()
_ = (a_jax @ b_jax).block_until_ready()
jax_uncompiled_time = time.time() - start
print(f"JAX (uncompiled): {jax_uncompiled_time:.3f}s")

# JAX JIT-compiled timing
def matmul(a, b):
    return a @ b

jitted_matmul = jit(matmul)
# First call compiles — don't time it
_ = jitted_matmul(a_jax, b_jax).block_until_ready()

# Now time the compiled version
start = time.time()
_ = jitted_matmul(a_jax, b_jax).block_until_ready()
jax_jit_time = time.time() - start
print(f"JAX (JIT): {jax_jit_time:.3f}s")

# Speedup ratio
print(f"\nSpeedup over NumPy: {numpy_time / jax_jit_time:.1f}x")

What does this tell us? NumPy is fast, but JAX with JIT is way faster — for large arrays. But if you run this same benchmark with a tiny 10×10 matrix, you’ll see JAX (even with JIT) is often slower than NumPy. We’ll come back to that trade-off later.

First, let’s understand what JAX actually is.

What Is JAX, Really? (And Why Should You Care?)

Here’s the most important thing to understand about JAX: it is not a drop-in NumPy replacement. It’s a function transformation library that happens to speak NumPy.

Think of it this way. NumPy is like a Swiss Army knife for arrays — it gives you all the tools you need, and you use them directly. JAX is more like a factory that can transform your tools automatically. Give JAX a function, and it can:

• Compile it — make it run faster on GPUs and TPUs (jax.jit)
• Differentiate it — compute exact derivatives automatically (jax.grad)
• Vectorize it — make it handle batches of data without rewriting (jax.vmap)
• Parallelize it — split it across multiple devices (jax.pmap)

The catch? For JAX to transform your function, the function must be pure — no side effects, no mutation. This is the hardest mental shift for NumPy users. In NumPy, you can write x[0] = 10 and the array changes in place. In JAX, that’s not allowed.

DeepMind’s adoption of JAX tells you this isn’t just a niche tool. As the DeepMind blog put it: JAX lets researchers ‘with no previous experience in high-performance computing… easily scale’ their code. That’s the promise.

Let’s see what JAX arrays look like.

# JAX arrays are not NumPy arrays
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
print(f"Type of JAX array: {type(x)}")
print(f"Contents: {x}")

# Compare to NumPy
import numpy as np
y = np.array([1.0, 2.0, 3.0])
print(f"Type of NumPy array: {type(y)}")

You’ll see jaxlib.xla_extension.ArrayImpl — a JAX-specific array type, not a NumPy ndarray. It behaves similarly in most ways (indexing, slicing, math), but there’s one huge difference.

Mind Shift #1: Arrays Don’t Change — You Recreate Them

Here’s the first gotcha that makes every NumPy user scream.

# This will raise an error
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
try:
    x[0] = 10  # ERROR: JAX arrays are immutable
except TypeError as e:
    print(f"Error: {e}")

You can’t mutate a JAX array. Instead, you use what looks like a bizarre syntax: x.at[0].set(10). This returns a new array with the update applied.

# The correct JAX way: use .at[].set()
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
x_updated = x.at[0].set(10.0)  # Returns a new array

print(f"Original: {x}")      # Still [1., 2., 3.]
print(f"Updated: {x_updated}")  # Now [10., 2., 3.]

‘But wait,’ you’re thinking — ‘that must be slow. Creating a new array every time you update a single element?’ It’s a fair concern. But here’s the thing: this immutability is what enables JAX’s speed.

When a function is pure (no side effects, no mutation), the JIT compiler can look at the whole computation, figure out what’s actually needed, and fuse operations together. It can eliminate intermediate arrays entirely. It can reorder operations for maximum parallelism. It can do all of this because it knows the function won’t change anything outside itself.

And the performance cost? It’s negligible at scale. Here’s a benchmark:

# Benchmark: in-place NumPy update vs JAX .at[].set() at different sizes
import numpy as np
import jax.numpy as jnp
import time

sizes = [1000, 1000000]  # 1K vs 1M elements

for n in sizes:
    print(f"\nArray size: {n:,} elements")
    
    # NumPy in-place update
    x_np = np.ones(n)
    start = time.time()
    for i in range(100):
        x_np[i % n] = i
    numpy_time = time.time() - start
    print(f"  NumPy in-place (100 updates): {numpy_time:.6f}s")
    
    # JAX .at[].set() — creates new array each time
    x_jax = jnp.ones(n)
    start = time.time()
    for i in range(100):
        x_jax = x_jax.at[i % n].set(float(i))
    # Force the last computation
    jax_time = time.time() - start
    print(f"  JAX .at[].set (100 updates): {jax_time:.6f}s")

At 1K elements, the difference is noticeable. At 1M elements, the overhead is buried by the cost of the actual computation. For machine learning workloads (where arrays are millions of elements), this isn’t a concern.

Mind Shift #2: Random Numbers Need an Explicit Key

Here’s another surprise for NumPy users. In NumPy, you set a seed and random numbers flow like magic:

import numpy as np
np.random.seed(42)
print(np.random.randn(3))  # [-0.280, -0.249, ...]

But that np.random.seed(42) is a side effect — it modifies hidden global state. JAX can’t allow that. Instead, every random function in JAX takes an explicit PRNG key.

# JAX random: explicit key management
import jax.random as jrandom

# Create a key
key = jrandom.PRNGKey(42)

# Generate three random numbers
random_values = jrandom.normal(key, shape=(3,))
print(f"First draw: {random_values}")

# If you use the same key again, you get the SAME numbers
random_values_again = jrandom.normal(key, shape=(3,))
print(f"Second draw (same key!): {random_values_again}")
print(f"Identical? {jnp.allclose(random_values, random_values_again)}")

This is a common beginner mistake — reusing the same key gives identical ‘random’ numbers. To get new random numbers, you need to split the key:

# Correct pattern: split the key to get new randomness
import jax.random as jrandom

key = jrandom.PRNGKey(42)

# Split: key becomes a pair (new_key, subkey)
key, subkey = jrandom.split(key)
random_values = jrandom.normal(subkey, shape=(3,))
print(f"New random values: {random_values}")

# Split again for another draw
key, subkey = jrandom.split(key)
more_random = jrandom.normal(subkey, shape=(3,))
print(f"More random values: {more_random}")

It feels annoying at first. But it’s actually better for reproducibility — every random number is traceable to a specific key, and there’s no hidden state that can accidentally be corrupted.

Why JAX Is Fast: JIT Compilation (and When It Isn’t)

We saw earlier that jax.jit can give a 15x speedup on large matrix multiplications. Now let’s understand what’s happening under the hood.

When you call jax.jit(f), JAX doesn’t immediately compile f. Instead, it traces the function — it runs it with abstract ‘tracer’ values that represent the shapes and types of your inputs, but not the actual data. From this trace, JAX builds a computation graph (called a jaxpr). Then it sends that graph to XLA, a compiler that optimizes and compiles it for your hardware (CPU, GPU, or TPU).

The result? The first call to a JIT-compiled function is slow (because it’s compiling). But all subsequent calls are fast — typically 5-10x faster than the uncompiled version for large arrays.

Let’s see this in action with the classic SeLU (Scaled Exponential Linear Unit) activation function:

# JIT benchmark: first call vs subsequent calls
import jax.numpy as jnp
from jax import jit
import time

# Define a moderately complex function (SeLU activation)
def selu(x, alpha=1.6732632423543772848170429916717, scale=1.0507009873554804934193349852946):
    return scale * jnp.where(x > 0, x, alpha * (jnp.exp(x) - 1))

# Create a large input
x = jnp.ones((5000, 5000))

# Uncompiled timing
start = time.time()
_ = selu(x).block_until_ready()
uncompiled_time = time.time() - start
print(f"Uncompiled: {uncompiled_time:.4f}s")

# JIT-compiled version
selu_jit = jit(selu)

# First call: includes compilation time
start = time.time()
_ = selu_jit(x).block_until_ready()
first_call_time = time.time() - start
print(f"JIT (first call, includes compilation): {first_call_time:.4f}s")

# Second call: steady-state performance
start = time.time()
_ = selu_jit(x).block_until_ready()
steady_time = time.time() - start
print(f"JIT (steady state): {steady_time:.4f}s")

print(f"\nSteady-state speedup: {uncompiled_time / steady_time:.1f}x")

But here’s the gotcha that trips everyone up. If your function has control flow that depends on values of arrays (like if x > 0 where x is an array), JIT tracing breaks:

# This will fail: control flow that depends on array values
from jax import jit
import jax.numpy as jnp

def bad_function(x):
    # This if statement depends on array values — JIT can't trace it
    if x > 0:
        return x
    else:
        return -x

try:
    jitted_bad = jit(bad_function)
    result = jitted_bad(3.0)  # ERROR!
except Exception as e:
    print(f"Error: {e}")
    print("\nFix: use jax.lax.cond")

To fix this, you use JAX’s control-flow operators like lax.cond (for if/else) or lax.switch (for switch/case):

# Fixed using jax.lax.cond
from jax import jit
import jax.lax as lax
import jax.numpy as jnp

def fixed_function(x):
    # lax.cond: condition, true_func, false_func, operand
    return lax.cond(x > 0, 
                    lambda x: x,   # true branch
                    lambda x: -x,  # false branch
                    x)

jitted_fixed = jit(fixed_function)
result = jitted_fixed(3.0)
print(f"Result: {result}")

This is the hardest part of using JAX’s JIT — you have to write your control flow in a way that JAX can trace. But once you learn the pattern, it becomes second nature.

Superpower #1: Automatic Differentiation with grad

Now we get to the really exciting part. jax.grad takes a Python function and returns a new function that computes its gradient — exact derivatives, automatically, with no manual formulas or numerical approximations.

Let’s start with the simplest possible example:

# Automatic differentiation with jax.grad
from jax import grad
import jax.numpy as jnp

# Define a simple function: f(x) = x^2
f = lambda x: x**2

# Get its gradient function
grad_f = grad(f)

# Evaluate the gradient at x = 3.0
result = grad_f(3.0)
print(f"f(x) = x^2")
print(f"f'(3.0) = {result}")  # Should be 6.0 (2 * 3)

The derivative of x^2 is 2x, so at x = 3.0, the gradient is 6.0. That’s exactly what we get.

Now let’s make it more realistic — a logistic regression model with multiple parameters:

# Gradient of a logistic regression loss function
from jax import grad
import jax.numpy as jnp

def sigmoid(x):
    return 1 / (1 + jnp.exp(-x))

def loss(weights, bias, x, y):
    """Binary cross-entropy loss for logistic regression."""
    logits = jnp.dot(x, weights) + bias
    preds = sigmoid(logits)
    # Binary cross-entropy
    return -jnp.mean(y * jnp.log(preds) + (1 - y) * jnp.log(1 - preds))

# Create synthetic data
key = jax.random.PRNGKey(0)
key, subkey = jax.random.split(key)

x = jax.random.normal(subkey, shape=(100, 5))  # 100 samples, 5 features
key, subkey = jax.random.split(key)
y = jax.random.randint(subkey, shape=(100,), minval=0, maxval=2).astype(jnp.float32)

weights = jnp.zeros(5)
bias = jnp.zeros(())

# Compute gradients with respect to both weights and bias
grad_loss = grad(loss, argnums=(0, 1))  # Differentiate w.r.t. weights (0) and bias (1)

grad_weights, grad_bias = grad_loss(weights, bias, x, y)

print(f"Gradient of weights shape: {grad_weights.shape}")
print(f"Gradient of weights: {grad_weights}")
print(f"Gradient of bias: {grad_bias}")

The argnums=(0, 1) tells grad which arguments to differentiate with respect to. This is incredibly powerful — you can compute gradients of complex models with nested loops and conditionals, as long as the function is pure and JIT-compilable.

And because grad returns a function, you can take higher-order derivatives — gradients of gradients:

# Higher-order derivatives: second derivative
from jax import grad

f = lambda x: x**3  # f'(x) = 3x^2, f''(x) = 6x

first_deriv = grad(f)
second_deriv = grad(first_deriv)  # gradient of gradient

print(f"f''(2.0) = {second_deriv(2.0)}")  # Should be 12.0 (6 * 2)

This is the foundation of most machine learning — gradient descent is just repeatedly evaluating grad(loss)(weights, data) and updating the weights.

Superpower #2: Automatic Vectorization with vmap

NumPy users spend a lot of time manually vectorizing code — reshaping arrays, adding batch dimensions, writing broadcast-compatible operations. JAX’s vmap automates this.

Here’s the idea: you write a function that processes a single example. Then vmap transforms it into a function that processes a batch of examples, automatically adding a batch dimension to all operations.

# Automatic vectorization with vmap
from jax import vmap
import jax.numpy as jnp
import time

# Function that processes a single example
def dot_product(a, b):
    """Dot product of two vectors."""
    return jnp.dot(a, b)

# Single pair
a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
print(f"Single dot product: {dot_product(a, b)}")

# Now we have 10000 pairs
a_batch = jnp.ones((10000, 3))
b_batch = jnp.ones((10000, 3))

# Option 1: Python loop (slow)
start = time.time()
results_loop = [dot_product(a_batch[i], b_batch[i]) for i in range(10000)]
loop_time = time.time() - start
print(f"Python loop: {loop_time:.4f}s")

# Option 2: vmap (fast, clean code)
vmap_dot = vmap(dot_product, in_axes=(0, 0))  # Batch along axis 0 for both inputs
start = time.time()
results_vmap = vmap_dot(a_batch, b_batch)
vmap_time = time.time() - start
print(f"vmap: {vmap_time:.4f}s")

print(f"vmap is {loop_time / vmap_time:.0f}x faster than Python loop")
print(f"Results match: {jnp.allclose(jnp.array(results_loop), results_vmap)}")

The in_axes=(0, 0) tells vmap which axis of each input represents the batch dimension (axis 0 for both). You can also specify out_axes for the output.

vmap is as fast as manual vectorization (both use XLA under the hood), but it’s much cleaner — you write simple, per-example logic and let JAX handle the batching.

Superpower #3: Multi-Device Parallelism with pmap

If you have multiple GPUs or a TPU, pmap is your friend. It replicates your function across all available devices, shards the data, and runs everything in parallel.

# Simple pmap example
from jax import pmap
import jax.numpy as jnp
from jax import local_device_count

# Create data: first dimension equals number of devices
n_devices = local_device_count()
data = jnp.arange(n_devices) ** 2

print(f"Number of devices: {n_devices}")
print(f"Input data: {data}")

# Define a function to run in parallel
def square(x):
    return x ** 2

# pmap replicates square across all devices
result = pmap(square)(data)
print(f"pmap result: {result}")

The key constraint: the mapped axis size must equal jax.local_device_count(). Each device processes one slice of the data.

pmap is the older API; for more complex parallelism, jax.shard_map (newer, more flexible) is now recommended. But pmap is still widely used for simple SPMD (Single Program, Multiple Data) workflows.

Putting It All Together: A Mini End-to-End Example

Let’s combine everything into a single workflow: training a linear regression model using gradient descent.

# End-to-end: linear regression with JAX
import jax.numpy as jnp
from jax import jit, grad, vmap
import jax.random as jrandom
import time

# Generate synthetic data
key = jrandom.PRNGKey(42)
key, subkey = jrandom.split(key)

# True function: y = 2x + 1 + noise
n_samples = 1000
x_data = jrandom.normal(subkey, shape=(n_samples, 1))
key, subkey = jrandom.split(key)
y_true = 2 * x_data + 1 + 0.1 * jrandom.normal(subkey, shape=(n_samples, 1))

# Model: linear regression
# predict = x @ weights + bias
def predict(params, x):
    weights, bias = params
    return jnp.dot(x, weights) + bias

def mse_loss(params, x, y):
    preds = predict(params, x)
    return jnp.mean((preds - y) ** 2)

# Gradient of loss with respect to (weights, bias)
grad_loss = grad(mse_loss)

# JIT-compiled update step
@jit
def update(params, x, y, lr=0.1):
    grads = grad_loss(params, x, y)
    weights, bias = params
    grad_weights, grad_bias = grads
    new_weights = weights - lr * grad_weights
    new_bias = bias - lr * grad_bias
    return (new_weights, new_bias)

# Initialize parameters
weights = jnp.zeros((1, 1))
bias = jnp.zeros(1)
params = (weights, bias)

# Training loop
n_epochs = 1000
start_time = time.time()
for epoch in range(n_epochs):
    params = update(params, x_data, y_true)
    if epoch % 200 == 0:
        loss_val = mse_loss(params, x_data, y_true)
        print(f"Epoch {epoch}: loss = {loss_val:.6f}")

training_time = time.time() - start_time
print(f"\nTraining completed in {training_time:.4f}s")
print(f"Learned weights: {params[0].item():.4f} (true: 2.0)")
print(f"Learned bias: {params[1].item():.4f} (true: 1.0)")

# Use vmap to predict on new data
new_x = jnp.array([[0.0], [1.0], [2.0]])
predictions = vmap(lambda x: predict(params, x))(new_x)
print(f"\nPredictions for x=[0, 1, 2]: {predictions.flatten()}")

Interpretation: After training, the weights should be close to 2.0 and the bias close to 1.0 (the true values we used to generate the data). The training time — even with 1000 epochs — is a fraction of a second, thanks to JIT compilation.

When NOT to Use JAX: Honest Trade-offs

JAX is not always the right tool. Here are the honest trade-offs:

• Small arrays (< 1000 elements): NumPy is faster. The dispatch overhead of JAX’s tracing and compilation doesn’t pay off for tiny computations. For a 10×10 matrix multiply, NumPy takes ~4ms while JAX takes ~170ms — that’s 42x slower.

• Debugging is harder. When a function is JIT-compiled, you lose Python stack traces. Errors come from XLA, not your code. Use jax.debug.print or disable JIT temporarily (with jax.disable_jit():) to debug.

• Smaller ecosystem. NumPy has SciPy, pandas, scikit-learn — JAX’s ecosystem is growing but smaller. Libraries like Equinox, Haiku, and Flax are filling the gap.

• Learning curve. Immutability, explicit PRNG keys, functional purity — these take time to internalize. If you just need simple array operations, NumPy is simpler.

Here’s the benchmark that shows the trade-off:

# Benchmark: NumPy vs JAX at different sizes
import numpy as np
import jax.numpy as jnp
from jax import jit
import time

sizes = [10, 100, 1000]

for n in sizes:
    print(f"\nMatrix size: {n}×{n}")
    
    # NumPy
    a_np = np.random.randn(n, n)
    b_np = np.random.randn(n, n)
    start = time.time()
    for _ in range(10):
        _ = a_np @ b_np
    numpy_time = (time.time() - start) / 10
    print(f"  NumPy: {numpy_time:.6f}s")
    
    # JAX with JIT
    a_jax = jnp.array(a_np)
    b_jax = jnp.array(b_np)
    
    @jit
    def matmul(a, b):
        return a @ b
    
    # Compile first
    _ = matmul(a_jax, b_jax).block_until_ready()
    
    start = time.time()
    for _ in range(10):
        _ = matmul(a_jax, b_jax).block_until_ready()
    jax_time = (time.time() - start) / 10
    print(f"  JAX JIT: {jax_time:.6f}s")
    
    if numpy_time < jax_time:
        print(f"  => NumPy {jax_time/numpy_time:.1f}x faster")
    else:
        print(f"  => JAX {numpy_time/jax_time:.1f}x faster")

At small sizes, NumPy wins. At large sizes, JAX wins. Know your problem size before choosing.

Recap: What You Learned

Let’s sum up what you now know about JAX:

1. Three mind shifts: Arrays are immutable (use .at[].set()), random numbers need explicit PRNG keys, and functions must be pure for transformations to work.

2. Three superpowers: jax.jit compiles your functions for speed (but watch out for control-flow gotchas), jax.grad gives you exact derivatives automatically, and jax.vmap vectorizes your code without manual batching.

3. Honest trade-offs: JAX is slower on small arrays, debugging is harder, and the ecosystem is smaller. But for large-scale computation, automatic differentiation, and multi-device parallelism, it’s transformative.

In the next part of this series, we’ll dive deeper into automatic differentiation — how grad actually works under the hood, and how to use it to build neural networks from scratch.

Check Your Understanding

Remember: List the three ‘mind shifts’ for NumPy users switching to JAX.

Understand: Why does JAX require functions to be pure (no side effects)? How does that enable JIT compilation?

Apply: Convert the following NumPy code to JAX:

import numpy as np
np.random.seed(0)
x = np.random.randn(5)
x[2] = 0

Analyze: Based on the benchmarks in this article, at what array size does JAX become faster than NumPy for matrix multiplication? What factors affect this threshold?

Evaluate: You’re building a data pipeline that processes 100-element arrays with simple statistics (mean, std, sum). Should you use JAX? Why or why not?

Create: Write a JAX function that computes the mean squared error between two arrays, then use jax.grad to get the gradient with respect to the first array. Verify your result by comparing to a finite-difference approximation.

Related Articles

• NumPy for Data Science: A Refresher (Part 1 of this series) — Covers the NumPy basics you need before switching to JAX.
• Introduction to Automatic Differentiation (upcoming Part 6) — Dives deeper into how grad works under the hood and how to build neural networks.
• JAX Quickstart: How to think in JAX — Official documentation, the best next step after this tutorial.