Run your first JAX program on NVIDIA GPUs with GKE

1. Introduction

JAX is a Python library for high-performance numerical computing. On the surface it looks like NumPy, but underneath it traces your Python functions, compiles them with XLA, and runs the result on accelerators such as NVIDIA GPUs.

In this codelab you provision a Google Kubernetes Engine cluster with NVIDIA L4 GPUs using Terraform, run JupyterLab inside the official NVIDIA JAX container on that GPU node, and write your first JAX computation. By the end you will have a working environment that the rest of this eight-part series builds on.

What you'll do

• Provision a GKE Standard cluster with a 2× NVIDIA L4 GPU node pool using Terraform
• Deploy JupyterLab on the GPU node from the NVIDIA JAX container image
• Verify the GPU end to end with nvidia-smi and jax.devices()
• Write JAX array code with jax.numpy and confirm the result lives on the GPU
• Apply the three transformations that define JAX: jax.jit, jax.grad, and jax.vmap

What you'll need

• A Google Cloud project with billing enabled, and workshop credits or a reservation covering GPU usage
• Quota for at least 2 NVIDIA L4 GPUs in your chosen region (how to check GPU quota)
• Basic Python and NumPy familiarity. No CUDA or Kubernetes experience required.

Estimated time to complete: 60 minutes.

2. Before you begin

Select your project

In the Google Cloud console, select or create a project with billing enabled.

Open Cloud Shell

Click Activate Cloud Shell (the terminal icon at the top right of the console) to start a Cloud Shell session, then point it at your project:

gcloud config set project <YOUR_PROJECT_ID>

Everything in this step runs in Cloud Shell, which already has gcloud, kubectl, terraform, and git installed.

Enable the required APIs

Enable every API this codelab needs in one command:

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

Confirm your GPU quota

The Terraform config requests two NVIDIA L4 GPUs. Confirm you have the quota:

gcloud compute regions describe us-central1 \
  --format="value(quotas.filter(metric:NVIDIA_L4_GPUS).limit)"

You should see a value of 2 or higher. If you see 0, request a quota increase before continuing.

Clone the workshop repository

The Terraform module and the Kubernetes manifest live in the workshop repository:

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu

The two directories you need are:

• terraform/ which contains GKE Standard cluster, VPC, node service account, and the L4 GPU node pool
• deploy/jupyter.yaml for PersistentVolumeClaim, JupyterLab Pod, and a LoadBalancer Service

3. Provision the GPU cluster with Terraform

A GPU node on GKE needs several things wired together: a VPC-native cluster, a node pool with an accelerator attached, and the NVIDIA driver installed on the node. The Terraform module does all three so you don't have to click through the console.

Configure your project

Copy the example variables file and point it at your project:

cd terraform
cp terraform.tfvars.example terraform.tfvars

Edit terraform.tfvars and set project_id. The defaults for everything else match this codelab:

project_id   = "<YOUR_PROJECT_ID>"
region       = "us-central1"
zone         = "us-central1-a"
cluster_name = "jax-gpu-cluster"
machine_type = "g2-standard-24"
gpu_type     = "nvidia-l4"
gpu_count    = 2

Understand what you are creating

Before you apply, look at the node pool definition in main.tf. This is the part that turns an ordinary node into a GPU node:

resource "google_container_node_pool" "gpu" {
  name     = "gpu-pool"
  location = var.zone
  cluster  = google_container_cluster.primary.name

  node_count = 1

  node_config {
    machine_type = var.machine_type

    guest_accelerator {
      type  = var.gpu_type   # nvidia-l4
      count = var.gpu_count  # 2

      gpu_driver_installation_config {
        gpu_driver_version = "DEFAULT"
      }
    }

    disk_size_gb = 100
    disk_type    = "pd-balanced"
    # ...
  }
}

Two details matter. First, machine_type and gpu_count must agree: g2-standard-24 comes with exactly 2 L4 GPUs, and g2-standard-48 comes with 4. Second, gpu_driver_installation_config is what makes the node usable — GKE installs the matching NVIDIA driver so your Pod only has to bring the CUDA user-space libraries.

Apply

terraform init
terraform apply

Review the plan and type yes. Cluster creation and node-pool provisioning take about 10 minutes. This is a good moment to read ahead.

When it finishes, fetch cluster credentials so kubectl talks to the new cluster:

$(terraform output -raw get_credentials_command)

Verify the node has GPUs

kubectl get nodes -o custom-columns=\
NAME:.metadata.name,GPU:.status.allocatable.nvidia\\.com/gpu

You should see output similar to:

NAME                                            GPU
gke-jax-gpu-cluster-gpu-pool-3f21a0b4-k7wq      2

4. Deploy JupyterLab on the GPU node

You now have a GPU node, but nothing running on it. The manifest in deploy/jupyter.yaml schedules a Pod that requests both GPUs and starts JupyterLab from the official NVIDIA JAX container.

Apply the manifest

cd ..
kubectl apply -f deploy/jupyter.yaml

This creates three objects:

• jax-workspace-pvc, a 50 GB persistent volume mounted at /workspace so your notebooks survive a Pod restart.
• jax-jupyter, the Pod that runs nvcr.io/nvidia/jax:26.04-maxtext-py3 and requests nvidia.com/gpu: "2".
• jax-jupyter-svc, a LoadBalancer that exposes JupyterLab on port 8884.

The GPU request is the important line:

resources:
  limits:
    nvidia.com/gpu: "2"
    memory: "48Gi"
    cpu: "12"

nvidia.com/gpu is an extended resource advertised by the GKE device plugin. Kubernetes only schedules this Pod onto a node that can satisfy it, which is how the Pod lands on your GPU node pool.

Wait for the Pod to be ready

The container image is large and the Pod also pip-installs JupyterLab on start, so the first pull takes several minutes:

kubectl get pod jax-jupyter -w

Wait until STATUS is Running, then press Ctrl+C.

Get the JupyterLab URL and token

Fetch the external IP of the Service:

kubectl get svc jax-jupyter-svc -w

Wait until EXTERNAL-IP changes from to an address, then press Ctrl+C.

JupyterLab prints a one-time login token to the Pod log:

kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

Open http://:8884 in your browser and paste the token when prompted.

Create a notebook

In JupyterLab, create a new Python 3 notebook in /workspace. Every code block in the rest of this codelab goes into a cell of that notebook.

5. Verify that JAX sees the GPU

Before writing any JAX code, confirm the hardware is visible from inside the container. If this step fails, nothing downstream will work.

Check the hardware

Run nvidia-smi from the notebook:

!nvidia-smi

You should see two L4 entries with a driver version and current memory usage.

Now ask for the compute capability, a two-digit number identifying the hardware generation. Later codelabs use features that depend on it: cuDNN fused attention needs 8.0 or newer, and FP8 needs 9.0 or newer.

import subprocess


def get_compute_capability() -> tuple[int, int]:
    """Query the compute capability of the first visible GPU."""
    out = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
        text=True,
    )
    major, minor = out.strip().split("\n")[0].split(".")
    return int(major), int(minor)


SM_MAJOR, SM_MINOR = get_compute_capability()
print(f"Detected compute capability: SM {SM_MAJOR}.{SM_MINOR}")

if SM_MAJOR < 7:
    print("WARNING: this course assumes SM 7.0+ (Volta or newer).")
else:
    print("GPU is compatible with this course.")

The L4 is an Ada Lovelace GPU, so you should see SM 8.9.

Check that JAX found the GPU

import jax
import jax.numpy as jnp

devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]

print(f"JAX version:        {jax.__version__}")
print(f"Default backend:    {jax.default_backend()}")
print(f"Available devices:  {devices}")

assert gpu_devices, f"No GPU backend found. Available devices: {devices}"
print(f"GPU devices:        {gpu_devices}")

You should see output similar to:

JAX version:        0.7.2
Default backend:    gpu
Available devices:  [CudaDevice(id=0), CudaDevice(id=1)]
GPU devices:        [CudaDevice(id=0), CudaDevice(id=1)]

The first import jax takes a few seconds because JAX initializes the CUDA runtime and probes devices.

How the pieces fit together

JAX itself never touches the GPU directly. It builds a program describing your computation and hands it down a stack:

Layer	Role
JAX	Traces your Python function into an intermediate representation
XLA	Compiles that IR into optimized GPU code
cuDNN, cuBLAS, NCCL	NVIDIA libraries XLA calls into for convolutions, GEMMs, and collectives
CUDA driver and runtime	Loads kernels onto the GPU and manages device memory

You almost never write CUDA yourself, but for most workloads the result is competitive with hand-written kernels. Codelab 2 opens up the tracing and compilation layer and codelab 3 shows you how to watch all of it execute in a profiler.

6. Write JAX array code on the GPU

The fastest way to get comfortable with JAX is to notice how much of it is just NumPy. The same constructors and broadcasting rules apply. What changes is where the array lives and how the computation runs.

Compare NumPy and JAX side by side

import numpy as np

# NumPy: runs on the CPU, stored in host memory
x_np = np.arange(8, dtype=np.float32)
y_np = np.sin(x_np) ** 2 + np.cos(x_np) ** 2
print(f"NumPy result:  {y_np}")
print(f"NumPy device:  CPU (host memory)")
print()

# JAX: same code, different array library
x = jnp.arange(8, dtype=jnp.float32)
y = jnp.sin(x) ** 2 + jnp.cos(x) ** 2
print(f"JAX result:    {y}")
print(f"JAX device:    {y.device}")

# Sanity check: the two answers should agree
np.testing.assert_allclose(y_np, np.asarray(y), atol=1e-6)
print()
print("NumPy and JAX agree.")

You should see output similar to:

JAX result:    [1. 1. 1. 1. 1. 1. 1. 1.]
JAX device:    cuda:0

Three things to notice:

1. The code is identical apart from np versus jnp.
2. y.device reports a CUDA device — JAX placed the array on the GPU automatically because that is the default backend.
3. JAX returns its own array type (jax.Array), not a NumPy array. Calling np.asarray(y) triggers a GPU to CPU transfer.

Watch that third point. Every crossing between GPU and host costs time, and printing a JAX array forces synchronization, because Python has to fetch the value to display it. Fine for a tiny example; a real problem inside a timed loop. The working rule is: create arrays with jnp, operate on them with jnp, and convert to NumPy only when you actually need to look at the values. Codelab 3 shows you how to spot accidental transfers in a profile.

7. Apply jit, grad, and vmap

NumPy on a GPU is useful, but on its own it is not a big leap over alternatives. What makes JAX distinctive is its function transformations: operators that take a Python function and return a new function with extra powers. Three of them show up in every remaining codelab.

jax.jit compiles your function

When you call a plain JAX function, operations dispatch to the GPU one at a time. Each dispatch has overhead, and small operations leave the GPU underused. jax.jit changes the execution model: JAX traces your function, XLA compiles it into an optimized executable, and JAX reuses that executable on later calls with compatible shapes and dtypes.

The first call is slow because it compiles. Every call after that is fast.

import time


def f(x):
    """Compose tanh, sin, and log1p so XLA has multiple ops to fuse when jitted."""
    return jnp.tanh(x) * jnp.sin(x) + jnp.log1p(x * x)


x = jnp.arange(1_000_000, dtype=jnp.float32)

# Eager: one kernel launch per operation
_ = f(x).block_until_ready()  # warm up
t0 = time.perf_counter()
for _ in range(10):
    y = f(x).block_until_ready()
eager_ms = (time.perf_counter() - t0) * 1000 / 10
print(f"Eager:            {eager_ms:6.3f} ms / call")

# Compiled: optimized executable, often with fused operations
f_jit = jax.jit(f)
_ = f_jit(x).block_until_ready()  # first call compiles
t0 = time.perf_counter()
for _ in range(10):
    y = f_jit(x).block_until_ready()
jit_ms = (time.perf_counter() - t0) * 1000 / 10
print(f"jax.jit (cached): {jit_ms:6.3f} ms / call")
print(f"Speedup:          {eager_ms / jit_ms:6.1f}x")

The exact speedup depends on the size and shape of your computation, but the pattern is universal: eager JAX is convenient, compiled JAX is fast.

jax.grad differentiates automatically

Training a neural network means computing gradients of a loss with respect to parameters. Pass any scalar-valued function to jax.grad and it returns a new function that computes the derivative.

def loss(w, x, y):
    """Mean squared error of `w*x` vs `y`; scalar loss for the `jax.grad` demo below."""
    pred = w * x
    return jnp.mean((pred - y) ** 2)


w = jnp.array(0.5)
xs = jnp.array([1.0, 2.0, 3.0, 4.0])
ys = jnp.array([2.0, 4.0, 6.0, 8.0])

# grad returns a function with the same signature, differentiating w.r.t. the first argument
dloss_dw = jax.grad(loss)

print(f"loss(w=0.5):   {loss(w, xs, ys):.4f}")
print(f"dloss/dw:      {dloss_dw(w, xs, ys):.4f}")

# Sanity check against a finite-difference approximation
eps = 1e-3
fd = (loss(w + eps, xs, ys) - loss(w - eps, xs, ys)) / (2 * eps)
print(f"finite diff:   {fd:.4f}  (should match)")

The gradient is negative, which tells an optimizer that increasing w will decrease the loss — exactly right, since the true relationship is y = 2x and you started at w = 0.5. Codelab 4 builds a full training loop around this.

jax.vmap vectorizes across a batch

GPUs like batched work. The naive way to apply a function to many inputs is a Python for loop, but that launches kernels one at a time and starves the GPU. jax.vmap takes a function written for a single example and returns a version that operates on a batch, with no loop and no manual reshaping.

def predict(W, x):
    """Tanh of a single-example matrix-vector product; vmapped below to batch over many `x`."""
    # Single example: W is (out, in), x is (in,) -> result is (out,)
    return jnp.tanh(W @ x)


key_w, key_x = jax.random.split(jax.random.key(0))
W = jax.random.normal(key_w, (4, 3))
xs = jax.random.normal(key_x, (10, 3))  # batch of 10 examples

# Without vmap: a Python loop, one kernel launch per example
ys_loop = jnp.stack([predict(W, x) for x in xs])

# With vmap: batch over the leading axis of xs, share W across the batch
batched_predict = jax.vmap(predict, in_axes=(None, 0))
ys_vmap = batched_predict(W, xs)

print(f"ys_loop shape:  {ys_loop.shape}")
print(f"ys_vmap shape:  {ys_vmap.shape}")
np.testing.assert_allclose(np.asarray(ys_loop), np.asarray(ys_vmap), atol=1e-6)
print("vmap matches the explicit loop.")

The in_axes=(None, 0) argument says: don't batch W (broadcast it), do batch xs along axis 0. The result is identical to the loop, but it dispatches as a single batched GPU operation.

Compose them

The real superpower is composition. Transformations stack:

fast_batched_grad = jax.jit(
    jax.vmap(jax.grad(loss), in_axes=(None, 0, 0))
)

One line gives you a compiled, vectorized, differentiated function that returns a per-example gradient for every (x, y) pair in a batch — most of what you need for batched training. Codelab 4 puts exactly this pattern to work on a real dataset.

8. Clean up

Delete the Jupyter workload, including the LoadBalancer and the persistent volume:

kubectl delete -f deploy/jupyter.yaml

Destroy the cluster, node pool, VPC, and service account:

cd terraform
terraform destroy

Type yes when prompted. Teardown takes a few minutes.

Finally, confirm nothing is left behind:

gcloud container clusters list
gcloud compute instances list

Both should be empty for this project. If you created a project just for this codelab, you can instead delete the whole project from the Cloud console.

9. Congratulations

You provisioned a GPU cluster from scratch and ran your first JAX program on it.

What you've learned

• How to provision a GKE Standard cluster with an NVIDIA L4 GPU node pool using Terraform, including the gpu_driver_installation_config block that makes the node usable
• How to schedule a Pod onto a GPU node with the nvidia.com/gpu extended resource
• How the JAX-on-GPU stack fits together: JAX traces, XLA compiles, and cuDNN, cuBLAS, and the CUDA runtime execute
• How to verify a GPU environment with nvidia-smi, compute capability, jax.devices(), and jax.default_backend()
• Where jax.numpy matches NumPy and where it differs: immutable arrays, .at[...] updates, default 32-bit dtypes, and the cost of host transfers
• How to apply and compose jax.jit, jax.grad, and jax.vmap

Next steps

• Codelab 2: Control JAX compilation with jax.jit where you will learn why the first call is slow, what triggers a recompile, and how to keep shapes stable
• Try JAX_PLATFORMS=cpu before importing JAX to force a CPU run, and compare the jax.jit timings above
• Scale the node pool to 4 L4 GPUs by setting machine_type = "g2-standard-48" and gpu_count = 4, and matching nvidia.com/gpu in deploy/jupyter.yaml

Reference docs

• JAX installation guide
• NVIDIA JAX container on NGC
• Run GPUs in GKE Standard node pools
• GPU regions and zones availability
• jax.jit, jax.grad, jax.vmap