I Built My First Gradient Descent Optimizer with JAX-Day 2

Aditya Goyal — Tue, 21 Jul 2026 08:43:20 +0000

On Day 1, I explored what JAX is and why people in machine learning love it. I learned about automatic differentiation ("grad") and Just-In-Time compilation ("jit"), but I hadn't actually used them in a real example.

Today I finally did.

The biggest "aha!" moment wasn't writing the code—it was understanding what gradients actually represent.

Thinking About Derivatives Differently

Imagine you're driving a car uphill.

You press the accelerator just a little. The car speeds up.

The question is:

How much does the speed change when you press the pedal slightly?

That's exactly what a derivative measures. It tells us how sensitive an output is to a tiny change in the input.

For example:

f(x) = x ** 3

At "x = 2":

f(2) = 8

Increase "x" a tiny bit:

f(2.001) ≈ 8.012

A very small increase in the input produces a measurable increase in the output.

Mathematically, the derivative of "x³" is:

3x²

So at "x = 2":

3 × 2² = 12

That means around this point, every tiny increase in "x" changes the output about 12 times as much.

So What Is a Gradient?

A derivative works when there's only one input.

But machine learning models don't have one input—they have thousands or even millions of parameters.

A gradient is simply the derivative with respect to every parameter.

One variable → Derivative

Many variables → Gradient (vector)

Instead of getting one number, we get a collection of numbers telling us how each parameter affects the loss.

Why Neural Networks Depend on Gradients

Every model starts with random parameters.

The model makes predictions.

Those predictions are compared with the correct answers using a loss function.

The gradient answers:

«"Which parameters should increase, and which should decrease, to reduce this loss?"»

Update the parameters.

Repeat.

Eventually the model learns.

That simple loop powers almost every deep learning model.

Trying Automatic Differentiation

Instead of computing derivatives by hand, JAX can do it automatically.

import jax

def cube(x):
return x ** 3

cube_grad = jax.grad(cube)

print(cube_grad(2.0))

Output:

12.0

Exactly what calculus predicts.

Creating a Tiny Learning Problem

Let's pretend we're trying to learn a simple equation:

prediction = weight × 4

The correct answer should be "20".

def loss(weight):
x = 4.0
target = 20.0

prediction = weight * x

return (prediction - target) ** 2

The ideal weight is:

because

5 × 4 = 20

Updating the Weight

grad_loss = jax.grad(loss)

weight = 0.5
lr = 0.05

for step in range(6):
gradient = grad_loss(weight)
weight -= lr * gradient

print(
    f"Step {step+1}: "
    f"weight={weight:.3f}, "
    f"loss={loss(weight):.3f}"
)

Example output:

Step 1: weight=4.100 loss=12.960
Step 2: weight=4.820 loss=0.518
Step 3: weight=4.964 loss=0.021
Step 4: weight=4.993 loss=0.001
Step 5: weight=4.999 loss=0.000
Step 6: weight=5.000 loss=0.000

Notice how the weight naturally moves closer to the correct answer after every iteration.

Making Computations Faster with "jit"

One of JAX's biggest strengths is compiling Python functions into highly optimized machine code.

import jax
import jax.numpy as jnp

def compute(x):
return jnp.sum(jnp.exp(-x) * jnp.sin(x))

compiled_compute = jax.jit(compute)

data = jnp.linspace(0, 100, 500000)

compiled_compute(data)

The first execution takes longer because JAX compiles the function.

After that, repeated calls become significantly faster.

This compilation overhead is usually worth it whenever the same computation runs many times during training.

My Biggest Takeaway

I used to think learning JAX meant memorizing functions.

Now I realize that's not the important part.

The real skill is understanding the ideas behind them.

"grad" computes derivatives automatically.
"jit" speeds up repeated computations.
Together, they form the foundation of efficient machine learning workflows.

I'm still early in my learning journey, but writing these examples myself made everything much clearer than simply reading documentation.

Next up: "vmap"—JAX's tool for automatically vectorizing computations. I'm excited to see how it can eliminate explicit Python loops.

DEV Community: Aditya Goyal

I Built My First Gradient Descent Optimizer with JAX-Day 2