{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Custom training loop from scratch in JAX",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/training-loop-in-jax",
  "publish_date": "2026-03-04",
  "corpus_stage": "",
  "notes": "",
  "id": "c8f222b2fc2b",
  "collected_at": "2026-10-02T14:11:03",
  "title": "Custom Training Loop from Scratch in JAX",
  "body_markdown": "Custom training loop from scratch in JAX\nWriting a training loop that you can actually control\nUpdated\nJuly 29, 2026\n•\n10\nmin read\n•\nView as Markdown\nW\nWesley Kambale\nI'm a Machine Learning Engineer passionate about building production-ready ML systems for the African market.\nWith experience in TensorFlow, Keras, and Python-based workflows, I help teams bridge the gap between machine learning research and real-world deployment—especially on resource-constrained devices.\nI'm also a Google Developer Expert in AI. I regularly speak at tech conferences including PyCon Africa, DevFest Kampala, DevFest Nairobi and more and also write technical articles on AI/ML here.\nPart of series\nJAX\nFor the past three weeks, we've been building up to this moment.\nWeek 1 taught us that JAX is fast. Week 2 showed us how to eliminate loops with\nvmap\nand compute gradients with\ngrad\n. Week 3 gave us Flax NNX—a way to define neural networks that feels like PyTorch but runs like JAX.\nBut we haven't actually\ntrained\nanything yet.\nToday, that changes. We're going to write a complete training loop from scratch: forward pass, loss calculation, gradient computation, parameter updates, and evaluation. No\nmodel.fit()\n. No magic. Just explicit, controllable, JIT-compiled code that you own completely.\nBy the end of this article, you'll have trained a CNN on MNIST and understood every single line of code that made it happen.\nWhy Write Your Own Training Loop?\nIf you've used Keras, you know how easy training can be:\nmodel.fit(x_train, y_train, epochs=10)\nOne line. Done. So why would anyone write hundreds of lines to do the same thing?\nBecause\nmodel.fit()\nis a black box. It works until it doesn't. And when it doesn't, when your loss explodes, when you need gradient accumulation, when you're debugging a custom layer, when you need to log something specific every 37 steps, you're stuck.\nThe training loop is where machine learning actually happens. If you don't control it, you don't control your model.\nIn production systems, you almost always need:\nCustom metrics\nthat aren't built into the framework\nGradient clipping\nto prevent exploding gradients\nLearning rate schedules\nthat change based on validation loss\nMixed precision training\nfor speed\nGradient accumulation\nwhen batches don't fit in memory\nEarly stopping\nbased on custom criteria\nWriting your own loop means you can implement any of these with a few lines of code.\nThe Four Beats of Training\nEvery training loop, regardless of framework, follows the same rhythm:\nForward Pass\n: Run input through the model to get predictions\nLoss Calculation\n: Compare predictions to ground truth\nBackward Pass\n: Compute gradients of the loss with respect to parameters\nParameter Update\n: Adjust parameters using the gradients\nIn PyTorch, this looks like:\noptimizer.zero_grad()        # Clear old gradients\npredictions = model(x)       # Forward pass\nloss = loss_fn(predictions, y)  # Loss calculation\nloss.backward()              # Backward pass\noptimizer.step()             # Parameter update\nPyTorch hides a lot here.\nloss.backward()\ntraverses a hidden computational graph and secretly updates\n.grad\nattributes on every parameter.\nJAX makes everything explicit.\nThe Optax Library\nOptax\nis JAX's optimization library. It handles gradient transformations and parameter updates.\nThe key insight about Optax is that optimizers are\ncomposable\n. Instead of monolithic optimizer classes, Optax gives you building blocks that you chain together:\nimport optax\n\n# A simple Adam optimizer\noptimizer = optax.adam(learning_rate=0.001)\n\n# Adam with gradient clipping\noptimizer = optax.chain(\n    optax.clip_by_global_norm(1.0),  # First, clip gradients\n    optax.adam(learning_rate=0.001)   # Then, apply Adam\n)\n\n# Adam with weight decay (AdamW)\noptimizer = optax.adamw(learning_rate=0.001, weight_decay=0.01)\nThis composability is powerful. Need to add gradient clipping to your existing optimizer? Add one line. Need a custom transformation? Write a function that takes gradients and returns modified gradients.\nCommon Optax Optimizers\nOptimizer\nUsage\noptax.sgd(lr)\nBasic stochastic gradient descent\noptax.adam(lr)\nAdam (most common default)\noptax.adamw(lr, weight_decay)\nAdam with decoupled weight decay\noptax.sgd(lr, momentum=0.9)\nSGD with momentum\noptax.rmsprop(lr)\nRMSprop\nCommon Gradient Transformations\nTransformation\nPurpose\noptax.clip_by_global_norm(max_norm)\nPrevent exploding gradients\noptax.add_decayed_weights(decay)\nL2 regularization\noptax.scale_by_schedule(schedule_fn)\nLearning rate scheduling\nThe nnx.Optimizer Wrapper\nOptax is a functional library, it works with raw JAX arrays. But our models are NNX objects with structured parameters. We need a bridge.\nnnx.Optimizer\nis that bridge:\nfrom flax import nnx\nimport optax\n\n# Create the Optax optimizer\ntx = optax.adam(learning_rate=0.001)\n\n# Wrap it with nnx.Optimizer\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\nThe\nwrt\nparameter is critical. It stands for \"with respect to\" and tells the optimizer which variables to update.\nnnx.Param\nmeans \"only update trainable parameters\", not batch normalization statistics or other state.\nOnce wrapped, updating parameters is simple:\noptimizer.update(model, grads)\nThis applies the gradients to the model's parameters in place. No need to manually extract parameters, apply updates, and put them back.\nComputing Gradients with nnx.value_and_grad\nIn PyTorch, you call\nloss.backward()\nto compute gradients. In JAX, we use\nvalue_and_grad\n.\nHere's the pattern:\ndef loss_fn(model, x, y):\n    logits = model(x)\n    loss = compute_loss(logits, y)\n    return loss\n\n# Get both the loss value and the gradients\nloss, grads = nnx.value_and_grad(loss_fn)(model, x, y)\nCritical rule\n:\nnnx.value_and_grad\ncomputes gradients with respect to the\nfirst argument\nof the function. That's why\nmodel\nmust be the first parameter.\nThe\ngrads\nobject has the same structure as the model's parameters—it's a pytree where each leaf is the gradient for the corresponding parameter.\nThe Complete Training Step\nLet's put it all together. Here's a complete, JIT-compiled training step:\n@nnx.jit\ndef train_step(model, optimizer, x, y):\n    \"\"\"Execute one training step.\"\"\"\n    \n    def loss_fn(model):\n        logits = model(x)\n        loss = optax.softmax_cross_entropy_with_integer_labels(\n            logits=logits, \n            labels=y\n        ).mean()\n        return loss, logits\n    \n    # Compute loss and gradients\n    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)\n    \n    # Update parameters\n    optimizer.update(model, grads)\n    \n    # Compute accuracy for logging\n    predictions = jnp.argmax(logits, axis=-1)\n    accuracy = jnp.mean(predictions == y)\n    \n    return loss, accuracy\nA few things to note:\n@nnx.jit\n: The entire function is JIT-compiled. The first call traces and compiles; subsequent calls execute the compiled code.\nhas_aux=True\n: Our loss function returns\n(loss, logits)\n. The\nhas_aux=True\nflag tells\nvalue_and_grad\nthat the function returns auxiliary data alongside the loss. Gradients are only computed for the first returned value.\nIn-place updates\n:\noptimizer.update(model, grads)\nmodifies the model in place. NNX handles the functional-to-stateful conversion internally.\nTraining a CNN on MNIST\nLet's train a real model on real data. We'll use the CNN from Week 3 and the MNIST dataset.\nSetup and Data Loading\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\nimport optax\nimport tensorflow_datasets as tfds\nimport tensorflow as tf\n\n# Disable TensorFlow GPU to avoid conflicts with JAX\ntf.config.set_visible_devices([], 'GPU')\n\n# Load MNIST\ndef load_mnist(batch_size=32):\n    \"\"\"Load and preprocess MNIST dataset.\"\"\"\n    \n    def preprocess(sample):\n        image = tf.cast(sample['image'], tf.float32) / 255.0\n        label = sample['label']\n        return {'image': image, 'label': label}\n    \n    train_ds = tfds.load('mnist', split='train')\n    test_ds = tfds.load('mnist', split='test')\n    \n    train_ds = (train_ds\n                .map(preprocess)\n                .shuffle(1024)\n                .batch(batch_size, drop_remainder=True)\n                .prefetch(1))\n    \n    test_ds = (test_ds\n               .map(preprocess)\n               .batch(batch_size, drop_remainder=True)\n               .prefetch(1))\n    \n    return train_ds, test_ds\n\ntrain_ds, test_ds = load_mnist(batch_size=32)\nDefine the Model\nWe'll use a slightly simplified version of the CNN from Week 3:\nfrom functools import partial\n\nclass CNN(nnx.Module):\n    \"\"\"Convolutional Neural Network for MNIST.\"\"\"\n    \n    def __init__(self, *, rngs: nnx.Rngs):\n        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n        self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n        self.pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n    \n    def __call__(self, x):\n        x = nnx.relu(self.conv1(x))\n        x = self.pool(x)\n        x = nnx.relu(self.conv2(x))\n        x = self.pool(x)\n        x = x.reshape(x.shape[0], -1)  # Flatten\n        x = nnx.relu(self.linear1(x))\n        x = self.linear2(x)\n        return x\nInitialize Model and Optimizer\n# Initialize model\nmodel = CNN(rngs=nnx.Rngs(0))\n\n# Test forward pass\ndummy_input = jnp.ones((1, 28, 28, 1))\ndummy_output = model(dummy_input)\nprint(f\"Model output shape: {dummy_output.shape}\")  # (1, 10)\n\n# Initialize optimizer\nlearning_rate = 0.005\nmomentum = 0.9\ntx = optax.adamw(learning_rate, momentum)\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n\n# Display model structure\nnnx.display(model)\nDefine Training and Evaluation Steps\n@nnx.jit\ndef train_step(model, optimizer, batch):\n    \"\"\"Single training step.\"\"\"\n    \n    def loss_fn(model):\n        logits = model(batch['image'])\n        loss = optax.softmax_cross_entropy_with_integer_labels(\n            logits=logits,\n            labels=batch['label']\n        ).mean()\n        return loss, logits\n    \n    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)\n    optimizer.update(model, grads)\n    \n    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == batch['label'])\n    return loss, accuracy\n\n@nnx.jit\ndef eval_step(model, batch):\n    \"\"\"Single evaluation step.\"\"\"\n    logits = model(batch['image'])\n    loss = optax.softmax_cross_entropy_with_integer_labels(\n        logits=logits,\n        labels=batch['label']\n    ).mean()\n    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == batch['label'])\n    return loss, accuracy\nThe Training Loop\n# Training configuration\nnum_epochs = 5\ntrain_steps_per_epoch = 1000\neval_steps = 200\n\nprint(\"Starting training...\")\nprint(\"=\" * 60)\n\nfor epoch in range(num_epochs):\n    # Training\n    model.train()  # Set model to training mode\n    train_loss, train_acc = 0.0, 0.0\n    \n    for step, batch in enumerate(train_ds.as_numpy_iterator()):\n        if step >= train_steps_per_epoch:\n            break\n        \n        # Convert to JAX arrays\n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        \n        loss, acc = train_step(model, optimizer, batch)\n        train_loss += loss\n        train_acc += acc\n    \n    train_loss /= train_steps_per_epoch\n    train_acc /= train_steps_per_epoch\n    \n    # Evaluation\n    model.eval()  # Set model to evaluation mode\n    eval_loss, eval_acc = 0.0, 0.0\n    eval_batches = 0\n    \n    for step, batch in enumerate(test_ds.as_numpy_iterator()):\n        if step >= eval_steps:\n            break\n        \n        batch = {k: jnp.array(v) for k, v in batch.items()}\n        \n        loss, acc = eval_step(model, batch)\n        eval_loss += loss\n        eval_acc += acc\n        eval_batches += 1\n    \n    eval_loss /= eval_batches\n    eval_acc /= eval_batches\n    \n    print(f\"Epoch {epoch + 1}/{num_epochs}\")\n    print(f\"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f}\")\n    print(f\"Eval Loss:  {eval_loss:.4f} | Eval Acc:  {eval_acc:.4f}\")\n\nprint(\"Training complete!\")\nUsing nnx.MultiMetric for Cleaner Tracking\nNNX provides a\nMultiMetric\nhelper for accumulating metrics across batches:\n# Define metrics\nmetrics = nnx.MultiMetric(\n    accuracy=nnx.metrics.Accuracy(),\n    loss=nnx.metrics.Average('loss'),\n)\n\n@nnx.jit\ndef train_step_with_metrics(model, optimizer, metrics, batch):\n    def loss_fn(model):\n        logits = model(batch['image'])\n        loss = optax.softmax_cross_entropy_with_integer_labels(\n            logits=logits, labels=batch['label']\n        ).mean()\n        return loss, logits\n    \n    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)\n    optimizer.update(model, grads)\n    \n    # Update metrics in place\n    metrics.update(loss=loss, logits=logits, labels=batch['label'])\n\n# In training loop:\nfor batch in train_ds:\n    train_step_with_metrics(model, optimizer, metrics, batch)\n\n# Get aggregated metrics\nresults = metrics.compute()\nprint(f\"Loss: {results['loss']:.4f}, Accuracy: {results['accuracy']:.4f}\")\n\n# Reset for next epoch\nmetrics.reset()\nThis is cleaner than manual accumulation and handles edge cases (like different batch sizes) correctly.\nInference: Using the Trained Model\nAfter training, making predictions is straightforward:\nmodel.eval()  # Ensure evaluation mode (deterministic)\n\n@nnx.jit\ndef predict(model, images):\n    logits = model(images)\n    return jnp.argmax(logits, axis=-1)\n\n# Get predictions\ntest_batch = next(test_ds.as_numpy_iterator())\npredictions = predict(model, jnp.array(test_batch['image']))\n\nprint(f\"Predictions: {predictions[:10]}\")\nprint(f\"Actual: {test_batch['label'][:10]}\")\nExercises\nAdd gradient clipping\n: Modify the optimizer to include\noptax.clip_by_global_norm(1.0)\n. Train the model and compare the loss curves.\nLearning rate schedule\n: Implement a learning rate that decays over time using\noptax.warmup_cosine_decay_schedule\n. Check the Optax documentation for the parameters.\nTrack more metrics\n: Add precision and recall tracking alongside accuracy. You'll need to compute true positives, false positives, and false negatives.\nEarly stopping\n: Modify the training loop to stop if validation accuracy doesn't improve for 3 consecutive epochs.\nQuick Reference\nimport optax\nfrom flax import nnx\n\n# Optimizer Setup\ntx = optax.adam(learning_rate=0.001)\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n\n# Training Step Pattern\n@nnx.jit\ndef train_step(model, optimizer, batch):\n    def loss_fn(model):\n        logits = model(batch['x'])\n        loss = compute_loss(logits, batch['y'])\n        return loss, logits\n    \n    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)\n    optimizer.update(model, grads)\n    return loss\n\n# Evaluation Step Pattern\n@nnx.jit  \ndef eval_step(model, batch):\n    logits = model(batch['x'])\n    loss = compute_loss(logits, batch['y'])\n    return loss\n\n# Training/Eval Mode\nmodel.train()  # Enable dropout, stochastic behavior\nmodel.eval()   # Deterministic inference\nWhat's Next\nWe've trained a model. But what happens when training goes wrong? What if loss becomes\nNaN\n? What if shapes don't match? What if the model silently produces garbage?\nNext week, we dive into\nOptax\nin more depth: learning rate schedules, gradient clipping, weight decay, and building custom optimizer stacks. We'll make our training more robust and production-ready.\n#\njax\n#\ntraining-loop\n#\nmachine-learning\n#\noptax",
  "metadata": {
    "title": "Custom Training Loop from Scratch in JAX",
    "description": "Write a complete training loop from scratch in JAX: forward pass, loss calculation, gradient computation, parameter updates, and evaluation.",
    "keywords": null,
    "author": null,
    "og:title": "Custom Training Loop from Scratch in JAX",
    "og:description": "Write a complete training loop from scratch in JAX: forward pass, loss calculation, gradient computation, parameter updates, and evaluation.",
    "og:url": "https://kambale.dev/training-loop-in-jax",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/46e5dfb7-7634-4ea1-9b04-a5b4909af30b.png",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:creator": "https://twitter.com/WesKambale",
    "twitter:title": "Custom Training Loop from Scratch in JAX",
    "twitter:description": "Write a complete training loop from scratch in JAX: forward pass, loss calculation, gradient computation, parameter updates, and evaluation.",
    "twitter:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/46e5dfb7-7634-4ea1-9b04-a5b4909af30b.png",
    "article:published_time": "2026-03-04T22:25:37.119Z",
    "article:modified_time": "2026-07-29T00:59:39.275Z",
    "article:author": "https://kambale.dev/@kambale"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/c8f222b2fc2b/code/body-000.txt",
      "chars": 39,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-001.txt",
      "chars": 239,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-002.txt",
      "chars": 358,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-003.txt",
      "chars": 183,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-004.txt",
      "chars": 31,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-005.txt",
      "chars": 199,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-006.txt",
      "chars": 637,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-007.txt",
      "chars": 993,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-008.txt",
      "chars": 764,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-009.txt",
      "chars": 406,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-010.txt",
      "chars": 870,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-011.txt",
      "chars": 1424,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-012.txt",
      "chars": 910,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-013.txt",
      "chars": 374,
      "language": "text"
    },
    {
      "file": "attachments/c8f222b2fc2b/code/body-014.txt",
      "chars": 759,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/46e5dfb7-7634-4ea1-9b04-a5b4909af30b.png",
      "bytes": 65642,
      "reused": true,
      "local_path": "attachments/c8f222b2fc2b/images/000.png",
      "src": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/46e5dfb7-7634-4ea1-9b04-a5b4909af30b.png",
      "alt": "Custom training loop from scratch in JAX",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/training-loop-in-jax.md",
      "label": "View as Markdown",
      "type": "a"
    },
    {
      "href": "https://hashnode.com/@kambale",
      "label": "W",
      "type": "a"
    },
    {
      "href": "https://twitter.com/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://linkedin.com/in/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/series/jax-series",
      "label": "Part of series JAX",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/training-loop",
      "label": "# training-loop",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/machine-learning",
      "label": "# machine-learning",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/optax",
      "label": "# optax",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:46Z",
  "source_url": "https://kambale.dev/training-loop-in-jax",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 15380,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/c8f222b2fc2b/pages/source.html",
  "local_html": "attachments/c8f222b2fc2b/pages/article.html",
  "body_base_path": "attachments/c8f222b2fc2b/pages/",
  "body_markdown_path": "attachments/c8f222b2fc2b/pages/article.md"
}