{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Optax: Optimizers You Can Compose Like LEGO",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/optax-optimizers",
  "publish_date": "2026-07-01",
  "corpus_stage": "",
  "notes": "",
  "id": "418e0493f3df",
  "collected_at": "2026-10-02T14:11:03",
  "title": "Optax: Optimizers You Can Compose Like LEGO",
  "body_markdown": "Optax: Optimizers You Can Compose Like LEGO\nUpdated\nJuly 29, 2026\n•\n11\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\nIn our\nlast article\non this blog in March, we built a complete training loop. We used\noptax.adamw()\nto create an optimizer, wrapped it in\nnnx.Optimizer\n, and watched our model learn.\nBut we barely scratched the surface.\nOptax isn't just a collection of optimizers. It's a\ncompositional system\nfor building gradient transformations. Instead of monolithic optimizer classes with dozens of parameters, Optax gives you small, focused building blocks that you chain together.\nWant Adam with gradient clipping? Chain them. Want SGD with momentum, weight decay, and a cosine learning rate schedule? Chain them. Want something completely custom that doesn't exist in any other framework? Build it from primitives.\nBy the end of this article, you'll understand:\nHow Optax's compositional design works\nHow to prevent exploding gradients with clipping\nHow to implement learning rate schedules without a separate scheduler object\nHow to apply different optimization strategies to different parameters\nHow to build production-grade optimizer stacks\nLet's take apart the LEGO box.\nThe Optax Mental Model\nIn PyTorch, an optimizer is an object that holds references to your parameters:\n# PyTorch\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)\noptimizer.step()  # Updates parameters in place\nIn Optax, an optimizer is a\npair of functions\n:\ninit(params)\n→ Creates the optimizer state (momentum buffers, etc.)\nupdate(grads, state, params)\n→ Returns\n(updates, new_state)\nYou then apply the updates to your parameters separately:\n# Pure Optax (without NNX)\nimport optax\n\noptimizer = optax.adam(learning_rate=0.001)\nopt_state = optimizer.init(params)\n\n# In training loop:\nupdates, opt_state = optimizer.update(grads, opt_state, params)\nnew_params = optax.apply_updates(params, updates)\nThis separation of concerns—computing updates vs. applying them—is what makes Optax so flexible. Each transformation only needs to know about gradients; it doesn't care where they came from or where they're going.\nWhen we use\nnnx.Optimizer\n, it handles the state management for us:\n# With NNX\noptimizer = nnx.Optimizer(model, optax.adam(0.001), wrt=nnx.Param)\noptimizer.update(model, grads)  # Handles everything internally\nBut understanding the underlying pattern helps when you need to do something custom.\nComposing Optimizers with optax.chain\nThe core of Optax's power is\noptax.chain()\n. It takes multiple gradient transformations and applies them in sequence:\nimport optax\n\n# Each transformation modifies the gradients, then passes them to the next\noptimizer = optax.chain(\n    optax.clip_by_global_norm(1.0),      # First: clip gradients\n    optax.scale_by_adam(),                # Second: compute Adam scaling\n    optax.add_decayed_weights(0.01),      # Third: add weight decay\n    optax.scale(-0.001)                   # Fourth: scale by -learning_rate\n)\nThe gradients flow through each transformation like water through pipes. Each pipe modifies the flow before passing it on.\nWhat optax.adam() Actually Is\nHere's something that surprises most people:\noptax.adam()\nisn't a primitive. It's a convenience wrapper around a chain:\n# optax.adam(learning_rate) is roughly equivalent to:\noptax.chain(\n    optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8),\n    optax.scale(-learning_rate)\n)\nAnd\noptax.adamw()\nadds weight decay:\n# optax.adamw(learning_rate, weight_decay) is roughly:\noptax.chain(\n    optax.scale_by_adam(),\n    optax.add_decayed_weights(weight_decay),\n    optax.scale(-learning_rate)\n)\nUnderstanding this lets you build exactly what you need.\nGradient Clipping: Preventing Explosions\nDeep networks, especially RNNs and Transformers, are prone to\nexploding gradients\n. One bad batch can produce gradients so large that a single update destroys your model.\nGradient clipping caps the magnitude of gradients before they're applied.\nGlobal Norm Clipping\nThe most common approach clips by the global norm—the total magnitude across all parameters:\noptimizer = optax.chain(\n    optax.clip_by_global_norm(max_norm=1.0),\n    optax.adam(learning_rate=0.001)\n)\nIf the global norm of your gradients exceeds 1.0, they're scaled down proportionally. If it's below 1.0, nothing happens.\nThis is what you want 90% of the time. It preserves the relative magnitudes between different parameters while preventing any single update from being catastrophically large.\nValue Clipping\nSometimes you want to clip each gradient element independently:\noptimizer = optax.chain(\n    optax.clip(max_delta=1.0),  # Clip each element to [-1, 1]\n    optax.adam(learning_rate=0.001)\n)\nThis is more aggressive and can distort the gradient direction, so use it carefully.\nMonitoring Gradient Norms\nBefore blindly adding clipping, it's useful to know what your gradient norms actually look like:\nimport jax\nimport jax.numpy as jnp\n\ndef compute_grad_norm(grads):\n    \"\"\"Compute the global L2 norm of gradients.\"\"\"\n    leaves = jax.tree_util.tree_leaves(grads)\n    return jnp.sqrt(sum(jnp.sum(g ** 2) for g in leaves))\n\n# In training:\nloss, grads = nnx.value_and_grad(loss_fn)(model)\ngrad_norm = compute_grad_norm(grads)\nprint(f\"Gradient norm: {grad_norm:.4f}\")\nIf you see norms of 100+ regularly, you need clipping. If they're consistently below 1, clipping won't change anything.\nLearning Rate Schedules\nIn PyTorch, learning rate scheduling requires a separate\nlr_scheduler\nobject that you call after each step or epoch:\n# PyTorch\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)\n\nfor batch in data:\n    optimizer.step()\n    scheduler.step()  # Don't forget this!\nIn Optax, the schedule is baked into the optimizer. No separate object, no extra\nstep()\ncall:\n# Optax\nschedule = optax.warmup_cosine_decay_schedule(\n    init_value=0.0,\n    peak_value=0.001,\n    warmup_steps=1000,\n    decay_steps=10000\n)\n\noptimizer = optax.adam(learning_rate=schedule)\nThe schedule is a function that takes a step count and returns a learning rate. Optax calls it automatically.\nCommon Schedules\nConstant (no schedule):\noptimizer = optax.adam(learning_rate=0.001)\nLinear warmup, then constant:\nschedule = optax.linear_schedule(\n    init_value=0.0,\n    end_value=0.001,\n    transition_steps=1000\n)\nWarmup, then cosine decay:\nschedule = optax.warmup_cosine_decay_schedule(\n    init_value=0.0,           # Start at 0\n    peak_value=0.001,         # Warm up to this\n    warmup_steps=1000,        # Over this many steps\n    decay_steps=9000,         # Then decay over this many steps\n    end_value=0.0001          # Down to this final value\n)\nExponential decay:\nschedule = optax.exponential_decay(\n    init_value=0.001,\n    transition_steps=1000,\n    decay_rate=0.96\n)\nPiecewise constant (step decay):\nschedule = optax.piecewise_constant_schedule(\n    init_value=0.001,\n    boundaries_and_scales={\n        5000: 0.1,   # At step 5000, multiply LR by 0.1\n        8000: 0.1,   # At step 8000, multiply by 0.1 again\n    }\n)\nVisualizing Schedules\nIt's helpful to plot your schedule before training:\nimport matplotlib.pyplot as plt\nimport jax.numpy as jnp\n\nschedule = optax.warmup_cosine_decay_schedule(\n    init_value=0.0,\n    peak_value=0.001,\n    warmup_steps=1000,\n    decay_steps=9000\n)\n\nsteps = jnp.arange(10000)\nlrs = [schedule(step) for step in steps]\n\nplt.figure(figsize=(10, 4))\nplt.plot(steps, lrs)\nplt.xlabel('Step')\nplt.ylabel('Learning Rate')\nplt.title('Warmup Cosine Decay Schedule')\nplt.grid(True)\nplt.show()\nWeight Decay: Regularization Done Right\nWeight decay prevents overfitting by penalizing large weights. But there's a subtle difference between L2 regularization and true weight decay.\nL2 regularization\nadds a term to the loss:\nloss = original_loss + λ * ||weights||²\nWeight decay\nsubtracts directly from the weights:\nweights = weights - lr * (grads + λ * weights)\nFor SGD, these are mathematically equivalent. For adaptive optimizers like Adam, they're not. The paper \"Decoupled Weight Decay Regularization\" showed that true weight decay works better with Adam—hence AdamW.\nIn Optax:\n# AdamW: Adam with decoupled weight decay\noptimizer = optax.adamw(learning_rate=0.001, weight_decay=0.01)\n\n# Or build it manually:\noptimizer = optax.chain(\n    optax.scale_by_adam(),\n    optax.add_decayed_weights(weight_decay=0.01),\n    optax.scale(-0.001)\n)\nExcluding Bias Terms from Weight Decay\nA common practice is to apply weight decay only to weight matrices, not to biases or layer norm parameters. Optax supports this with masks:\nfrom flax import nnx\nimport optax\n\n# Create a mask that's True for kernel/weight parameters, False for biases\ndef create_weight_decay_mask(params):\n    def should_decay(path, _):\n        # path is a tuple like ('linear1', 'kernel') or ('linear1', 'bias')\n        return 'kernel' in path or 'weight' in path\n    \n    return jax.tree_util.tree_map_with_path(\n        lambda path, _: should_decay(path, _),\n        params\n    )\n\n# Use optax.masked to apply weight decay selectively\noptimizer = optax.chain(\n    optax.clip_by_global_norm(1.0),\n    optax.adamw(learning_rate=0.001, weight_decay=0.01, mask=create_weight_decay_mask)\n)\nPer-Parameter Optimization\nSometimes you want different learning rates for different parts of your model. This is common when fine-tuning: you might want a small learning rate for pretrained layers and a larger one for the new classification head.\nOptax handles this with\noptax.multi_transform\n:\nimport optax\nfrom flax import nnx\n\n# Define different optimizers for different parameter groups\noptimizer = optax.multi_transform(\n    transforms={\n        'backbone': optax.adam(learning_rate=1e-5),   # Small LR for pretrained\n        'head': optax.adam(learning_rate=1e-3),       # Larger LR for new layers\n    },\n    param_labels=param_labels  # A pytree matching params, with string labels\n)\nThe tricky part is creating\nparam_labels\n—a pytree with the same structure as your parameters, where each leaf is a string label.\nHere's a practical example:\nfrom flax import nnx\nimport jax\n\nclass FineTuneModel(nnx.Module):\n    def __init__(self, *, rngs: nnx.Rngs):\n        # Pretrained backbone (we want small LR)\n        self.backbone = nnx.Linear(784, 256, rngs=rngs)\n        \n        # New classification head (we want larger LR)\n        self.head = nnx.Linear(256, 10, rngs=rngs)\n    \n    def __call__(self, x):\n        x = nnx.relu(self.backbone(x))\n        return self.head(x)\n\nmodel = FineTuneModel(rngs=nnx.Rngs(0))\n\n# Create parameter labels\ndef label_fn(path, _):\n    \"\"\"Assign labels based on parameter path.\"\"\"\n    path_str = '/'.join(str(p) for p in path)\n    if 'backbone' in path_str:\n        return 'backbone'\n    else:\n        return 'head'\n\n# Get params and create labels\nparams = nnx.state(model, nnx.Param)\nparam_labels = jax.tree_util.tree_map_with_path(label_fn, params)\n\n# Create the multi-transform optimizer\ntx = optax.multi_transform(\n    transforms={\n        'backbone': optax.adam(learning_rate=1e-5),\n        'head': optax.adam(learning_rate=1e-3),\n    },\n    param_labels=param_labels\n)\n\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\nBuilding a Production Optimizer Stack\nLet's put everything together into a production-grade optimizer configuration:\nimport optax\n\ndef create_optimizer(\n    peak_learning_rate: float = 1e-3,\n    warmup_steps: int = 1000,\n    total_steps: int = 100000,\n    weight_decay: float = 0.01,\n    max_grad_norm: float = 1.0,\n    end_learning_rate: float = 1e-5,\n):\n    \"\"\"Create a production-ready optimizer with all the bells and whistles.\"\"\"\n    \n    # Learning rate schedule: warmup then cosine decay\n    schedule = optax.warmup_cosine_decay_schedule(\n        init_value=0.0,\n        peak_value=peak_learning_rate,\n        warmup_steps=warmup_steps,\n        decay_steps=total_steps - warmup_steps,\n        end_value=end_learning_rate,\n    )\n    \n    # Build the optimizer chain\n    optimizer = optax.chain(\n        # 1. Gradient clipping (prevent explosions)\n        optax.clip_by_global_norm(max_grad_norm),\n        \n        # 2. Adam scaling\n        optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8),\n        \n        # 3. Weight decay (applied after Adam scaling)\n        optax.add_decayed_weights(weight_decay),\n        \n        # 4. Scale by negative learning rate (makes it gradient descent)\n        optax.scale_by_schedule(lambda step: -schedule(step)),\n    )\n    \n    return optimizer\n\n# Usage\ntx = create_optimizer(\n    peak_learning_rate=1e-3,\n    warmup_steps=1000,\n    total_steps=50000,\n    weight_decay=0.01,\n    max_grad_norm=1.0,\n)\n\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\nThis gives you:\nGradient clipping to prevent explosions\nAdam's adaptive learning rates\nDecoupled weight decay for regularization\nWarmup to stabilize early training\nCosine decay to fine-tune convergence\nThis is essentially what's used to train modern Transformers.\nInspecting Optimizer State\nSometimes you need to peek inside the optimizer. With NNX:\n# Get the full optimizer state\nopt_state = nnx.state(optimizer)\n\n# The state is a nested structure containing:\n# - Step count\n# - Momentum buffers (for Adam's first moment)\n# - Velocity buffers (for Adam's second moment)\n# - Any other state from transformations in the chain\n\nprint(jax.tree_util.tree_map(lambda x: x.shape, opt_state))\nThis is useful for:\nDebugging optimization issues\nImplementing custom logging\nUnderstanding memory usage (optimizer state can be 2-3x the model size!)\nCommon Optimizer Recipes\nHere are ready-to-use configurations for common scenarios:\nStandard Training (Most Cases)\ntx = optax.adamw(learning_rate=1e-3, weight_decay=0.01)\nTransformer Training\nschedule = optax.warmup_cosine_decay_schedule(\n    init_value=0.0, peak_value=1e-4, warmup_steps=2000, decay_steps=98000\n)\ntx = optax.chain(\n    optax.clip_by_global_norm(1.0),\n    optax.adamw(learning_rate=schedule, weight_decay=0.1)\n)\nFine-Tuning (Small LR, No Warmup)\ntx = optax.adamw(learning_rate=2e-5, weight_decay=0.01)\nWhen Training is Unstable\ntx = optax.chain(\n    optax.clip_by_global_norm(0.5),  # Aggressive clipping\n    optax.sgd(learning_rate=0.01, momentum=0.9)  # SGD is more stable\n)\nExercises\nVisualize your schedule\n: Create a warmup + cosine decay schedule and plot it. Experiment with different warmup lengths.\nCompare clipping strategies\n: Train the same model with\nclip_by_global_norm(1.0)\n,\nclip_by_global_norm(0.1)\n, and no clipping. Monitor the gradient norms and final accuracy.\nImplement a custom schedule\n: Create a schedule that uses a high learning rate for the first 1000 steps, then drops to 1/10th, then decays linearly to zero.\nFine-tuning setup\n: Take a model with two parts (backbone + head) and set up a multi_transform optimizer with 10x different learning rates.\nQuick Reference\nimport optax\n\n# Basic Optimizers\noptax.sgd(learning_rate, momentum=0.9)\noptax.adam(learning_rate)\noptax.adamw(learning_rate, weight_decay)\n\n# Gradient Clipping\noptax.clip_by_global_norm(max_norm)\noptax.clip(max_delta)\n\n# Learning Rate Schedules\noptax.warmup_cosine_decay_schedule(init, peak, warmup_steps, decay_steps)\noptax.exponential_decay(init, transition_steps, decay_rate)\noptax.piecewise_constant_schedule(init, boundaries_and_scales)\n\n# Composing\noptax.chain(transform1, transform2, transform3)\n\n# Building Blocks\noptax.scale_by_adam()           # Adam's moment estimation\noptax.add_decayed_weights(wd)   # Weight decay\noptax.scale(-lr)                # Scale by learning rate\noptax.scale_by_schedule(fn)     # Dynamic scaling\n\n# With NNX\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\noptimizer.update(model, grads)\nWhat's Next\nWe've built training loops and sophisticated optimizers. But what happens when things go wrong? What if your loss becomes NaN? What if shapes don't match?\nNow that we are back, next week, we'll tackle\nreliability engineering with Chex\n. We'll learn to write assertions that catch bugs before they corrupt your model, validate tensor shapes at compile time, and hunt down the source of NaN values.\nBecause fast code that produces garbage is still garbage.\nNext week\n:\nCatching Bugs Before They Catch You; Reliability Engineering with Chex\n#\njax\n#\noptimizers",
  "metadata": {
    "title": "Optax: Optimizers You Can Compose Like LEGO",
    "description": "In our last article on this blog in March, we built a complete training loop. We used optax.adamw() to create an optimizer, wrapped it in nnx.Optimizer, and watched our model learn.\nBut we barely scra",
    "keywords": null,
    "author": null,
    "og:title": "Optax: Optimizers You Can Compose Like LEGO",
    "og:description": "In our last article on this blog in March, we built a complete training loop. We used optax.adamw() to create an optimizer, wrapped it in nnx.Optimizer, and watched our model learn.\nBut we barely scra",
    "og:url": "https://kambale.dev/optax-optimizers",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/5f81618f-bb0f-4c8f-b781-23307cb393fd.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": "Optax: Optimizers You Can Compose Like LEGO",
    "twitter:description": "In our last article on this blog in March, we built a complete training loop. We used optax.adamw() to create an optimizer, wrapped it in nnx.Optimizer, and watched our model learn.\nBut we barely scra",
    "twitter:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/5f81618f-bb0f-4c8f-b781-23307cb393fd.png",
    "article:published_time": "2026-07-01T08:27:02.850Z",
    "article:modified_time": "2026-07-29T01:00:03.751Z",
    "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/418e0493f3df/code/body-000.txt",
      "chars": 117,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-001.txt",
      "chars": 255,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-002.txt",
      "chars": 142,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-003.txt",
      "chars": 398,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-004.txt",
      "chars": 154,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-005.txt",
      "chars": 174,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-006.txt",
      "chars": 108,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-007.txt",
      "chars": 126,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-008.txt",
      "chars": 364,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-009.txt",
      "chars": 231,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-010.txt",
      "chars": 192,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-011.txt",
      "chars": 44,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-012.txt",
      "chars": 103,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-013.txt",
      "chars": 314,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-014.txt",
      "chars": 107,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-015.txt",
      "chars": 219,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-016.txt",
      "chars": 425,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-017.txt",
      "chars": 40,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-018.txt",
      "chars": 47,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-019.txt",
      "chars": 259,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-020.txt",
      "chars": 629,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-021.txt",
      "chars": 396,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-022.txt",
      "chars": 1114,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-023.txt",
      "chars": 1376,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-024.txt",
      "chars": 336,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-025.txt",
      "chars": 56,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-026.txt",
      "chars": 237,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-027.txt",
      "chars": 56,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-028.txt",
      "chars": 149,
      "language": "text"
    },
    {
      "file": "attachments/418e0493f3df/code/body-029.txt",
      "chars": 830,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/5f81618f-bb0f-4c8f-b781-23307cb393fd.png",
      "bytes": 67302,
      "reused": true,
      "local_path": "attachments/418e0493f3df/images/000.png",
      "src": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/5f81618f-bb0f-4c8f-b781-23307cb393fd.png",
      "alt": "Optax: Optimizers You Can Compose Like LEGO",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/optax-optimizers.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/training-loop-in-jax",
      "label": "last article",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/optimizers",
      "label": "# optimizers",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:47Z",
  "source_url": "https://kambale.dev/optax-optimizers",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 16733,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/418e0493f3df/pages/source.html",
  "local_html": "attachments/418e0493f3df/pages/article.html",
  "body_base_path": "attachments/418e0493f3df/pages/",
  "body_markdown_path": "attachments/418e0493f3df/pages/article.md"
}