JAX from Scratch: Step-by-Step Guide to

What is JAX?

JAX is a high-performance numerical computing library from Google Research that combines the familiar NumPy API with automatic differentiation and accelerated hardware execution. At its core, JAX provides composable function transformations that let you write numerical code once and automatically derive gradients, vectorize across batches, and compile to XLA-optimized kernels for CPUs, GPUs, and TPUs.

Unlike frameworks that require you to learn an entirely new programming model, JAX is designed to feel like writing NumPy or native Python — but with the ability to tap into the full power of modern hardware accelerators. The key insight is that JAX treats Python functions as first-class objects that can be transformed by higher-order functions.

Core Transformations

JAX provides four primary transformations that work on pure Python functions:

• grad — Automatic differentiation for computing gradients
• jit — Just-in-time compilation via XLA for performance
• vmap — Automatic vectorization for batching operations
• pmap — Parallel mapping across multiple devices

These transformations can be freely composed. You can take the gradient of a jitted function, vectorize a differentiated function, or compile a batched gradient computation — all without modifying the original source code.

Why JAX Matters

JAX addresses fundamental pain points in scientific computing and machine learning research:

1. NumPy Compatibility with Acceleration

JAX provides jax.numpy (typically imported as jnp) which mirrors the NumPy API almost perfectly. This means existing NumPy code can often be ported to JAX by simply changing the import statement — and immediately gain GPU acceleration.

2. Automatic Differentiation Without Graph Building

Unlike TensorFlow's static graph approach or PyTorch's tape-based autodiff, JAX uses source-code-transformation-based automatic differentiation. Functions are pure Python, and gradients are computed by transforming the function's code directly. This yields clean, readable derivative code and supports higher-order differentiation naturally — you can compute gradients of gradients (Hessians) by simply calling grad(grad(f)).

3. Functional Programming Paradigm

JAX encourages pure functional programming: functions should not have side effects, and all data should be immutable. This constraint enables the powerful transformation system but also leads to cleaner, more testable, and more parallelizable code.

4. Research to Production Continuity

Because JAX uses XLA (the same compiler infrastructure as TensorFlow), models developed in research settings can be deployed to production without rewriting. The same JAX code runs efficiently on a laptop CPU, a multi-GPU workstation, or a TPU pod.

Getting Started: Installation and Setup

Install JAX using pip. Choose the appropriate version based on your hardware:

# CPU-only version
pip install jax jaxlib

# GPU version (CUDA 12)
pip install jax jaxlib[cuda12]

# TPU version (Google Cloud)
pip install jax jaxlib[tpu]

After installation, verify your setup and check available devices:

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

print(f"JAX version: {jax.__version__}")
print(f"Available devices: {jax.devices()}")
print(f"Default device: {jax.devices()[0]}")

# Quick sanity check: GPU-accelerated matrix multiplication
key = jax.random.PRNGKey(42)
a = jax.random.normal(key, (1000, 1000))
b = jax.random.normal(key, (1000, 1000))
c = jnp.dot(a, b)
print(f"Result shape: {c.shape}, device: {c.device()}")

JAX Fundamentals: Step by Step

Step 1: jax.numpy — The Accelerated NumPy

The entry point to JAX is jax.numpy. It looks and feels like NumPy but with crucial differences rooted in JAX's functional nature. Let's explore the basics:

import jax.numpy as jnp
from jax import random

# Creating arrays
x = jnp.array([1.0, 2.0, 3.0, 4.0])
y = jnp.ones((3, 4))
z = jnp.arange(0, 10, 0.5)

# Random number generation (requires explicit PRNG keys)
key = random.PRNGKey(0)
random_array = random.normal(key, shape=(100,))

# Element-wise operations
result = jnp.sin(x) + jnp.cos(y)

# Linear algebra
matrix = jnp.array([[1., 2.], [3., 4.]])
eigenvalues, eigenvectors = jnp.linalg.eigh(matrix)
print("Eigenvalues:", eigenvalues)

Critical difference: JAX arrays are immutable. You cannot do in-place updates like x[i] = 5.0. Instead, use index-based updates with the functional syntax:

# Instead of x[0] = 99.0, do:
x_updated = x.at[0].set(99.0)

# For multiple indices
x_updated = x.at[1:3].set(jnp.array([10.0, 20.0]))

# Conditional updates
x_updated = x.at[x > 2.0].set(0.0)

# Additive updates
x_updated = x.at[0].add(5.0)

Step 2: Understanding PRNG Keys

JAX uses a functional random number generation system based on splittable PRNG keys. This is one of the most important concepts to master early:

from jax import random

# Create an initial key
key = random.PRNGKey(42)

# Split the key for independent random streams
key1, key2 = random.split(key)
key3, key4 = random.split(key2)

# Use keys to generate random numbers
samples1 = random.normal(key1, shape=(1000,))
samples2 = random.uniform(key3, shape=(500,), minval=-1.0, maxval=1.0)

# Keys are consumed — generate new ones after use
key1, subkey = random.split(key1)
more_samples = random.normal(subkey, shape=(500,))

The key principle: never reuse a PRNG key. Each key should be used exactly once for random number generation, then split to obtain fresh keys. This ensures reproducibility and mathematical correctness in stochastic computations.

Step 3: Automatic Differentiation with grad

JAX's automatic differentiation is accessed through the grad transformation. It takes a function and returns a new function that computes the gradient:

from jax import grad, value_and_grad
import jax.numpy as jnp

# Define a scalar-valued function
def f(x):
    return jnp.sum(x ** 2)

# Get the gradient function
grad_f = grad(f)

# Evaluate gradient
x = jnp.array([1.0, 2.0, 3.0])
gradient = grad_f(x)
print("Gradient of sum(x²) at [1,2,3]:", gradient)  # Should be [2, 4, 6]

# For functions with multiple arguments, use argnums
def loss(params, x, y):
    predictions = params @ x
    return jnp.sum((predictions - y) ** 2)

# Gradient with respect to first argument (params)
grad_loss_wrt_params = grad(loss, argnums=0)

# Gradient with respect to multiple arguments
grad_loss_wrt_params_and_x = grad(loss, argnums=(0, 1))

To compute both the function value and its gradient in one pass (avoiding redundant computation), use value_and_grad:

# Compute value and gradient simultaneously
value_and_grad_f = value_and_grad(f)
value, gradient = value_and_grad_f(x)
print(f"Value: {value}, Gradient: {gradient}")

Step 4: Higher-Order Derivatives

JAX makes higher-order differentiation trivial by composing grad:

from jax import grad, jacfwd, jacrev, hessian

def f(x):
    return jnp.sum(x ** 3)

# First derivative
first_derivative = grad(f)

# Second derivative (Hessian-vector product style)
second_derivative = grad(grad(f))

# Direct Hessian computation
hess = hessian(f)

x = jnp.array([1.0, 2.0])
print("Gradient:", first_derivative(x))
print("Hessian diagonal:", second_derivative(x))
print("Full Hessian:", hess(x))

For vector-valued functions, use jacfwd (forward-mode) or jacrev (reverse-mode) depending on the input/output dimensions:

# Vector-valued function
def g(x):
    return jnp.stack([x[0]**2 + x[1], x[0] * x[1]**2])

# Reverse-mode Jacobian (efficient for many outputs)
jacobian_rev = jacrev(g)

# Forward-mode Jacobian (efficient for few inputs)
jacobian_fwd = jacfwd(g)

x = jnp.array([2.0, 3.0])
print("Jacobian (reverse mode):", jacobian_rev(x))

Step 5: JIT Compilation with jit

The jit transformation compiles functions using XLA, yielding dramatic speedups for repeated computations:

from jax import jit
import time

# Define a computationally intensive function
def slow_function(x):
    for _ in range(100):
        x = jnp.sin(x) + jnp.cos(x) * jnp.tanh(x)
    return x

# JIT-compiled version
fast_function = jit(slow_function)

# First call triggers compilation (slower)
x = jnp.ones((1000, 1000))
start = time.time()
result1 = fast_function(x)
result1.block_until_ready()  # Ensure computation is done
print(f"First call (includes compilation): {time.time() - start:.3f}s")

# Subsequent calls use cached compilation (much faster)
start = time.time()
result2 = fast_function(x)
result2.block_until_ready()
print(f"Second call (cached): {time.time() - start:.3f}s")

You can also use jit as a decorator:

@jit
def neural_network_layer(x, weights, bias):
    return jnp.tanh(x @ weights + bias)

# This function is now compiled automatically
x = jax.random.normal(jax.random.PRNGKey(0), (128, 64))
w = jax.random.normal(jax.random.PRNGKey(1), (64, 32))
b = jax.random.normal(jax.random.PRNGKey(2), (32,))
output = neural_network_layer(x, w, b)

Step 6: Tracing and Static Arguments

JIT works by tracing your function with abstract values. This means Python control flow must be resolvable at trace time. Use static_argnums for arguments that determine control flow:

from jax import jit
import jax.numpy as jnp

# This will fail: x depends on runtime value
@jit
def bad_poly(x, n):
    result = 0.0
    for i in range(n):  # n determines loop structure
        result = result + x ** i
    return result

# Correct: mark n as static
@jit
def good_poly(x, n):
    result = jnp.zeros_like(x)
    for i in range(n):
        result = result + x ** i
    return result

# Use static_argnums for the jit call
poly = jit(good_poly, static_argnums=(1,))

x = jnp.array([0.5, 0.5])
print(poly(x, 5))   # Works: n=5 is static
print(poly(x, 10))  # Works but triggers recompilation for n=10

Step 7: Automatic Vectorization with vmap

The vmap transformation lifts a function to operate over batches automatically, eliminating manual batch dimension handling:

from jax import vmap
import jax.numpy as jnp

# A function that operates on a single sample
def predict_single(params, x):
    w, b = params
    return jnp.dot(w, x) + b

# Vectorize it to handle batches
predict_batch = vmap(predict_single, in_axes=(None, 0))

# params is not batched (None), x is batched along axis 0
params = (jnp.array([1.0, 2.0, 3.0]), jnp.array(0.5))
x_batch = jnp.array([[1.0, 1.0, 1.0],
                      [2.0, 2.0, 2.0],
                      [3.0, 3.0, 3.0]])

predictions = predict_batch(params, x_batch)
print("Batch predictions:", predictions)

# vmap can also map over multiple axes
def kernel(x, y):
    return jnp.exp(-jnp.sum((x - y) ** 2))

# Compute all pairwise kernels
pairwise_kernel = vmap(vmap(kernel, in_axes=(None, 0)), in_axes=(0, None))
# Equivalent to: vmap(vmap(kernel, (None, 0)), (0, None))

vmap is especially powerful when combined with grad:

# Compute per-example gradients efficiently
def per_example_loss(params, x, y):
    w, b = params
    pred = jnp.dot(w, x) + b
    return (pred - y) ** 2

# Vectorize the gradient computation
per_example_grad = vmap(grad(per_example_loss, argnums=0), in_axes=(None, 0, 0))

params = (jnp.array([1.0, 2.0]), jnp.array(0.5))
x_batch = jnp.array([[1.0, 1.0], [2.0, 3.0], [4.0, 5.0]])
y_batch = jnp.array([2.0, 3.0, 4.0])

gradients = per_example_grad(params, x_batch, y_batch)
print("Per-example gradients:", gradients)

Step 8: Composing Transformations

The true power of JAX emerges when you compose transformations. Here's a complete example combining grad, vmap, and jit:

from jax import grad, vmap, jit, value_and_grad
import jax.numpy as jnp

# Define a simple neural network layer
def layer(params, x):
    w, b = params
    return jnp.tanh(x @ w + b)

# Define a two-layer MLP
def mlp(params, x):
    w1, b1, w2, b2 = params
    h = jnp.tanh(x @ w1 + b1)
    return h @ w2 + b2

# Mean squared error loss
def mse_loss(params, x, y):
    pred = mlp(params, x)
    return jnp.mean((pred - y) ** 2)

# Compose: jit + vmap + value_and_grad for batch training
@jit
def batch_loss_and_grad(params, x_batch, y_batch):
    # Vectorize loss computation over batch
    batch_losses = vmap(lambda x, y: mse_loss(params, x, y))(x_batch, y_batch)
    return jnp.mean(batch_losses)

# Get gradient function
grad_fn = jit(grad(mse_loss, argnums=0))

# Full training step with composed transformations
@jit
def training_step(params, x_batch, y_batch, learning_rate=0.01):
    loss, grads = value_and_grad(mse_loss, argnums=0)(params, x_batch, y_batch)
    # Update parameters (functional style)
    w1, b1, w2, b2 = params
    dw1, db1, dw2, db2 = grads
    new_params = (
        w1 - learning_rate * dw1,
        b1 - learning_rate * db1,
        w2 - learning_rate * dw2,
        b2 - learning_rate * db2
    )
    return new_params, loss

Building a Complete Training Pipeline

Let's put everything together into a realistic MNIST classifier training loop. This example demonstrates the end-to-end workflow:

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap, value_and_grad, random
import numpy as np
from functools import partial

# 1. Data loading (simulated MNIST-like data)
def generate_data(key, num_samples=1000, input_dim=784, num_classes=10):
    k1, k2, k3 = random.split(key, 3)
    X = random.normal(k1, (num_samples, input_dim))
    true_weights = random.normal(k2, (input_dim, num_classes))
    logits = X @ true_weights
    labels = jnp.argmax(logits, axis=-1)
    return X, labels

# 2. Model definition
def init_params(key, layer_sizes):
    k1, k2, k3, k4 = random.split(key, 4)
    w1 = random.normal(k1, (layer_sizes[0], layer_sizes[1])) * 0.1
    b1 = jnp.zeros(layer_sizes[1])
    w2 = random.normal(k2, (layer_sizes[1], layer_sizes[2])) * 0.1
    b2 = jnp.zeros(layer_sizes[2])
    return (w1, b1, w2, b2)

def forward(params, x):
    w1, b1, w2, b2 = params
    h = jnp.tanh(x @ w1 + b1)
    logits = h @ w2 + b2
    return logits

# 3. Loss function
def cross_entropy_loss(params, x, y_onehot):
    logits = forward(params, x)
    log_probs = jax.nn.log_softmax(logits)
    return -jnp.sum(log_probs * y_onehot) / len(y_onehot)

# 4. Accuracy metric
def accuracy(params, x_batch, y_batch):
    logits = vmap(partial(forward, params))(x_batch)
    predictions = jnp.argmax(logits, axis=-1)
    return jnp.mean(predictions == y_batch)

# 5. Training step (composed transformations)
@jit
def train_step(params, opt_state, x_batch, y_batch, y_onehot, learning_rate=0.001):
    loss, grads = value_and_grad(cross_entropy_loss, argnums=0)(params, x_batch, y_onehot)
    # Simple SGD update
    w1, b1, w2, b2 = params
    dw1, db1, dw2, db2 = grads
    new_params = (
        w1 - learning_rate * dw1,
        b1 - learning_rate * db1,
        w2 - learning_rate * dw2,
        b2 - learning_rate * db2
    )
    return new_params, loss

# 6. Full training loop
key = random.PRNGKey(42)
k_data, k_init = random.split(key)

# Generate synthetic data
X_train, y_train = generate_data(k_data, num_samples=5000)
X_test, y_test = generate_data(random.split(k_data)[0], num_samples=1000)
y_train_onehot = jax.nn.one_hot(y_train, 10)
y_test_onehot = jax.nn.one_hot(y_test, 10)

# Initialize model
params = init_params(k_init, [784, 128, 10])

# Training configuration
batch_size = 64
num_epochs = 20
num_batches = len(X_train) // batch_size

# Training loop
for epoch in range(num_epochs):
    epoch_loss = 0.0
    # Shuffle data each epoch
    perm_key, key = random.split(key)
    perm = random.permutation(perm_key, len(X_train))
    X_shuffled = X_train[perm]
    y_shuffled = y_train_onehot[perm]
    
    for batch_idx in range(num_batches):
        start = batch_idx * batch_size
        end = start + batch_size
        x_batch = X_shuffled[start:end]
        y_batch = y_train[start:end]
        y_onehot_batch = y_shuffled[start:end]
        
        params, loss = train_step(params, None, x_batch, y_batch, y_onehot_batch)
        epoch_loss += loss
    
    # Evaluate
    train_acc = accuracy(params, X_train[:1000], y_train[:1000])
    test_acc = accuracy(params, X_test, y_test)
    print(f"Epoch {epoch+1}: Loss={epoch_loss/num_batches:.4f}, "
          f"Train Acc={train_acc:.3f}, Test Acc={test_acc:.3f}")

print("\nTraining complete!")
print(f"Final test accuracy: {accuracy(params, X_test, y_test):.4f}")

Advanced Patterns and Best Practices

1. PyTrees for Structured Parameters

JAX natively understands PyTrees — structures of arrays like dicts, lists, tuples, and custom classes. Instead of manually tracking parameter tuples, use dictionaries or dataclasses:

from dataclasses import dataclass
from typing import Any
import jax.tree_util as jtu

@dataclass
class ModelParams:
    encoder: Any  # dict or nested structure
    decoder: Any
    embedding: jnp.ndarray

# JAX can automatically map over this structure
params = ModelParams(
    encoder={"w": jnp.ones((64, 32)), "b": jnp.zeros(32)},
    decoder={"w": jnp.ones((32, 64)), "b": jnp.zeros(64)},
    embedding=jnp.ones((1000, 64))
)

# grad works seamlessly with PyTrees
def loss_fn(params, x):
    # Access nested parameters naturally
    h = x @ params.encoder["w"] + params.encoder["b"]
    return jnp.sum(h)

grads = grad(loss_fn)(params, jnp.ones((16, 64)))
print(type(grads))  # ModelParams with same structure
print(grads.encoder["w"].shape)  # (64, 32)

2. Using lax for Low-Level Control

When jax.numpy isn't sufficient, jax.lax provides primitive operations with finer control:

from jax import lax
import jax.numpy as jnp

# Custom scan (loop with state)
def cumulative_sum(x):
    def body(carry, xi):
        new_carry = carry + xi
        return new_carry, new_carry
    return lax.scan(body, 0.0, x)[1]

# Custom while loop
def newton_sqrt(a, tolerance=1e-6):
    def cond(state):
        x, _ = state
        return jnp.abs(x * x - a) > tolerance
    
    def body(state):
        x, _ = state
        x = 0.5 * (x + a / x)
        return x, x
    
    init_x = a / 2.0
    final_x, _ = lax.while_loop(cond, body, (init_x, init_x))
    return final_x

print(newton_sqrt(25.0))  # Approximately 5.0

# Custom conditionals with lax.cond
def relu_with_leaky(x, negative_slope=0.01):
    return lax.cond(
        x > 0,
        lambda: x,
        lambda: negative_slope * x
    )

3. Handling Randomness in jit Functions

When using JIT-compiled functions that need randomness, thread PRNG keys through the computation:

from jax import jit, random

@jit
def dropout_layer(key, x, drop_rate=0.5):
    kept_prob = 1.0 - drop_rate
    mask = random.bernoulli(key, p=kept_prob, shape=x.shape)
    return mask * x / kept_prob

# Thread keys through a training loop
def training_epoch(params, key, data):
    losses = []
    for x_batch in data:
        key, subkey = random.split(key)
        loss = compute_loss_with_dropout(params, subkey, x_batch)
        losses.append(loss)
    return jnp.mean(jnp.array(losses)), key  # Return updated key

4. Memory Management and Asynchronous Dispatch

JAX dispatches operations asynchronously to devices. Use block_until_ready() for timing benchmarks and when you need guaranteed completion:

import time

@jit
def heavy_computation(x):
    return jnp.sin(x) @ jnp.cos(x).T

x = jnp.ones((2048, 2048))

# Without blocking — may measure dispatch time only
result = heavy_computation(x)

# With blocking — measures actual computation
result = heavy_computation(x)
result.block_until_ready()

# Proper benchmarking
start = time.perf_counter()
result = heavy_computation(x)
result.block_until_ready()
elapsed = time.perf_counter() - start
print(f"Actual computation time: {elapsed:.4f}s")

5. Donations for In-Place Updates

While JAX arrays are immutable, you can use buffer donation to reuse memory when you know an input won't be needed after the function call:

@jit
def update_params(params, grads):
    # Standard functional update — allocates new arrays
    return jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)

# With donation (advanced)
@jit
def update_params_inplace(params, grads):
    # Mark params for donation after use
    return jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)

# Usage requires explicit donate_argnums
params = jnp.ones((1000, 1000))
grads = jnp.ones((1000, 1000)) * 0.1
new_params = jit(update_params_inplace, donate_argnums=(0,))(params, grads)

Common Pitfalls and How to Avoid Them

Pitfall 1: In-Place Updates

Problem: Using NumPy-style in-place mutation.

# This raises an error in JAX
x = jnp.array([1.0, 2.0, 3.0])
# x[0] = 99.0  # TypeError: JAX arrays are immutable

Solution: Use .at updates or rewrite with pure functional operations.

x_updated = x.at[0].set(99.0)

Pitfall 2: Python Control Flow in jit

Problem: JIT traces with abstract values, so Python conditionals on array values fail.

@jit
def buggy_relu(x):
    if x > 0:  # x is a traced array, not a concrete value
        return x
    else:
        return 0.0

Solution: Use jnp.where, lax.cond, or mark problematic arguments as static.

@jit
def correct_relu(x):
    return jnp.where(x > 0, x, 0.0)

Pitfall 3: Reusing PRNG Keys

Problem: Reusing a key produces identical "random" numbers, breaking stochastic algorithms.

key = random.PRNGKey(0)
x = random.normal(key, (100,))
y = random.normal(key, (100,))  # x == y exactly!

Solution: Always split keys before use.

key = random.PRNGKey(0)
k1, k2 = random.split(key)
x = random.normal(k1, (100,))
y = random.normal(k2, (100,))  # Independent random numbers

Pitfall 4: Large JIT Compilation Overheads

Problem: JIT compilation on first call can be slow, and recompilation happens when static arguments or shapes change.

Solution: Use static_argnums sparingly, pad inputs to fixed shapes, and pre-compile critical functions:

# Pre-compile by calling once with dummy data
dummy_input = jnp.ones((128, 784))
_ = my_jitted_model(dummy_input)  # Triggers compilation

# Now use with real data — no compilation overhead
real_output = my_jitted_model(real_input)

Pitfall 5: Ignoring Device Placement

Problem: Arrays default to the first device, causing unnecessary transfers.

Solution: Explicitly place arrays and use jax.device_put:

# Explicit device placement
cpu_device = jax.devices("cpu")[0]
gpu_device = jax.devices("gpu")[0] if jax.devices("gpu") else cpu_device

x_on_gpu = jax.device_put(jnp.ones((1000, 1000)), gpu_device)
# Now computations involving x_on_gpu stay on GPU

Real-World Example: Training a CNN from Scratch

Here's a complete convolutional neural network implementation demonstrating JAX best practices with PyTrees, vmap, and composed transformations:

import jax
import jax.numpy as jnp
from jax import random, grad, jit, vmap, value_and_grad, lax
import numpy as np

# Define CNN primitives using lax
def conv2d(x, kernel, stride=1):
    """Simple 2D convolution with padding"""
    k_h, k_w, in_c, out_c = kernel.shape
    batch, h, w, _ = x.shape
    pad_h, pad_w = k_h // 2, k_w // 2
    x_padded = jnp.pad(x, ((0, 0), (pad_h, pad_h), (pad_w, pad_w), (0, 0)))
    
    # Extract patches using lax.conv_general_dilated
    dimension_numbers = lax.ConvDimensionNumbers(
        lhs_spec=(0, 3, 1, 2),  # (N, C, H, W)
        rhs_spec=(3, 2, 0, 1),  # (O, I, H, W)
        out_spec=(0, 3, 1, 2)   # (N, C, H, W)
    )
    return lax.conv_general_dilated(
        x_padded.transpose(0, 3, 1, 2),
        kernel.transpose(3, 2, 0, 1),
        window_strides=(stride, stride),
        padding="SAME",
        dimension_numbers=dimension_numbers
    ).transpose(0, 2, 3, 1)

# Initialize CNN parameters as a PyTree dictionary
def init_cnn(key):
    k1, k2, k3, k4 = random.split(key, 4)
    return {
        "conv1_kernel": random.normal(k1, (3, 3, 1, 16)) * 0.1,
        "conv1_bias": jnp.zeros(16),
        "conv2_kernel": random.normal(k2, (3, 3, 16, 32)) * 0.1,
        "conv2_bias": jnp.zeros(32),
        "fc_weight": random.normal(k3, (32 * 7 * 7, 10)) * 0.01,
        "fc_bias": jnp.zeros(10),
    }

def cnn_forward(params, x):
    # Conv1 + ReLU + Pool
    h = conv2d(x, params["conv1_kernel"]) + params["conv1_bias"]
    h = jnp.relu(h)
    h = h[:, ::2, ::2, :]  # Max pool 2x2 (simplified as stride pool)
    
    # Conv2 + ReLU + Pool
    h = conv2d(h, params["conv2_kernel"]) + params["conv2_bias"]
    h = jnp.relu(h)
    h = h[:, ::2, ::2, :]  # Now shape: (batch, 7, 7, 32)
    
    # Flatten + FC
    h = h.reshape(h.shape[0], -1)
    logits = h @ params["fc_weight"] + params["fc_bias"]
    return