{"slug": "i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2", "title": "I Built a Neural Network Training Loop in 5 Lines Using JAX - Day 2", "summary": "A developer built a neural network training loop in five lines using JAX, demonstrating gradient descent from scratch and achieving a 32x speedup with JIT compilation on a MacBook CPU. The project verifies JAX's gradient computation by hand and explains the core concepts of derivatives and gradients for machine learning.", "body_md": "SERIES: Learning RL and JAX in Public - from zero to DeepMind :)\n\nOn Day 1 I covered what JAX is and the three superpowers: GPU acceleration, grad, and jit. Today I actually used them. And the moment where I manually verified a gradient by hand and watched JAX confirm it is the one that made everything real for me.\n\nBut before the code, I had to understand what a gradient actually is. Because I realised I had been using them for years without a clean mental model.\n\n**What is a derivative, really?**\n\nImagine you are adjusting a dial on a machine. The dial controls the volume. Your goal is to get the volume to exactly 50. Right now it is at 30.\n\nYou turn the dial a tiny bit. The volume goes up.\n\nThe derivative answers one question: if I nudge this dial by a tiny amount, how much does the output change?\n\nThat is it. In math:\n\n```\nf(x) = x * x\nf(3) = 9\nf(3.001) = 9.006001\n\nnudge of 0.001 caused a change of 0.006\nrate of change = 0.006 / 0.001 = 6\n```\n\nThe derivative of x squared at x equals 3 is 6. Meaning: at this point, if x increases by a tiny amount, the output increases 6 times as fast.\n\n**What is a gradient?**\n\nA derivative is for a function with one input. A gradient is the same idea for a function with many inputs. It is just a list of derivatives, one per input.\n\n``` php\nfunction with 1 input  ->  derivative  (a single number)\nfunction with N inputs ->  gradient    (a list of N numbers)\n```\n\n**Why does this matter for ML?**\n\nEvery neural network has weights. Training means finding the weights that make the model least wrong. \"How wrong\" is measured by a number called the loss.\n\nThe gradient of the loss tells you: for each weight, does increasing it make the loss bigger or smaller, and by how much?\n\nNudge every weight in the direction that reduces the loss. Repeat thousands of times. That is gradient descent. That is literally how every neural network on the planet trains.\n\n**Now the code**\n\nFirst I verified grad by hand:\n\n``` python\nimport jax\n\ndef square(x):\n    return x * x\n\ngrad_of_square = jax.grad(square)\nprint(grad_of_square(3.0))  # 6.0\n```\n\nDerivative of x squared is 2x. At x equals 3, that is 6. JAX says 6.0. Verified.\n\nThen I wrote a real loss function:\n\n``` python\ndef loss(w):\n    x = 2.0\n    target = 10.0\n    prediction = w * x\n    return (prediction - target) ** 2\n\ngrad_of_loss = jax.grad(loss)\n```\n\nThis loss measures how wrong a prediction is. The ideal weight is 5.0 because 5 times 2 equals 10.\n\nThen gradient descent from scratch:\n\n```\nlearning_rate = 0.1\nw = 1.0\n\nfor step in range(5):\n    grad = grad_of_loss(w)\n    w = w - learning_rate * grad\n    print(f\"step {step + 1}: w = {w:.4f}, loss = {loss(w):.4f}\")\n```\n\nOutput:\n\n```\nstep 1: w = 2.6000, loss = 21.1600\nstep 2: w = 3.6640, loss = 8.6491\nstep 3: w = 4.3302, loss = 3.5369\nstep 4: w = 4.7481, loss = 1.4461\nstep 5: w = 5.0105, loss = 0.0044 (almost zero, w approaching 5.0)\n```\n\nThat loop is the core of every neural network training run. More weights, more complex functions, but the same mechanic.\n\n**The jit speedup**\n\nI benchmarked a large computation with and without jit over 100 runs:\n\n``` python\ndef slow_computation(x):\n    return jnp.sum(jnp.sin(x) ** 2 + jnp.cos(x) ** 2)\n\nfast_computation = jax.jit(slow_computation)\n```\n\nResult: 32x speedup on my MacBook CPU. On a GPU this goes to 100x or more.\n\nThe first call is always slower because that is when JAX compiles. Every call after uses the cached compiled version. One catch: if you change the shape of your input, JAX recompiles. So keep your array shapes consistent.\n\n**One question I had: do I need to memorise all this syntax?**\n\nNo. The important thing is knowing what to use and when. Nobody in a research role writes jax.grad from memory every time - they know it exists, they know what it does, they look up the exact syntax when needed.\n\nWhat you do need to carry in your head: grad differentiates, jit compiles, and they compose together. That mental model handles 90% of situations.\n\nDay 3 tomorrow: vmap, the fourth JAX superpower. This one genuinely surprised me.\n\nAll code from this series, organised by day, is on my GitHub: [https://github.com/MadhumithaKolkar/jax-rl-lab](https://github.com/MadhumithaKolkar/jax-rl-lab)\n\nHappy learning everyone !", "url": "https://wpnews.pro/news/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2", "canonical_source": "https://dev.to/madhumithakolkar/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2-1a4j", "published_at": "2026-07-21 04:24:25+00:00", "updated_at": "2026-07-21 04:59:06.417216+00:00", "lang": "en", "topics": ["machine-learning", "neural-networks", "developer-tools"], "entities": ["JAX", "DeepMind"], "alternates": {"html": "https://wpnews.pro/news/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2", "markdown": "https://wpnews.pro/news/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2.md", "text": "https://wpnews.pro/news/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2.txt", "jsonld": "https://wpnews.pro/news/i-built-a-neural-network-training-loop-in-5-lines-using-jax-day-2.jsonld"}}