{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Never lose your progress: Checkpointing with Orbax",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/orbax",
  "publish_date": "2026-08-07",
  "corpus_stage": "",
  "notes": "",
  "id": "ab459134afd4",
  "collected_at": "2026-10-02T14:10:59",
  "title": "Never lose your progress: Checkpointing with Orbax",
  "body_markdown": "Never lose your progress: Checkpointing with Orbax\nUpdated\nAugust 7, 2026\n•\n16\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\nPicture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes out. Or the Colab runtime disconnects. Or your cloud instance gets preempted because you are on a spot pricing tier to save money. Six hours of compute, gone, with nothing to show for it but a chart that no longer exists anywhere except your memory of watching it.\nThis is not a hypothetical for anyone building AI systems outside of a well funded lab with guaranteed, uninterrupted access to a GPU cluster. Power instability is a genuine constraint in many parts of the world. Free tier compute gets reclaimed without warning. Long training runs get interrupted by things entirely outside your control.\nThe solution is checkpointing: periodically saving the complete state of your training run to disk, so that if something goes wrong, you can restore exactly where you left off instead of starting over from scratch.\nToday we cover\nOrbax\n, the checkpointing library built for the JAX ecosystem. Orbax does for training state what Git does for code: it gives you reliable, versioned snapshots you can always return to. By the end of this article you will understand:\nWhy checkpointing in JAX requires a different approach than\ntorch.save()\nHow Flax NNX exposes model state as something Orbax can save\nHow to save and restore model parameters using\nCheckpointManager\nHow to save multiple pieces of state together, such as both model and optimizer\nHow checkpointing works when your model is sharded across multiple devices\nHow to build a complete, crash resilient training loop\nLet's make sure we never lose six hours of compute again.\nWhy checkpointing in JAX is different\nIn PyTorch, saving a model is famously simple:\ntorch.save(model.state_dict(), \"checkpoint.pt\")\nThis works because\nstate_dict()\nreturns a Python dictionary of tensors, and PyTorch's serialization just pickles that dictionary to disk. Loading is equally simple: build the model architecture, then call\nmodel.load_state_dict(torch.load(\"checkpoint.pt\"))\n.\nJAX, as we have discussed since\nweek 1\n, has no hidden state. There is no single object holding \"the model\" that you can wave a wand over and serialize. What exists instead are pytrees, explicit, nested structures of arrays that your training loop passes around by hand.\nThis has real advantages once you understand it. Because JAX state is always an explicit, structured pytree, checkpointing can be done with the same rigor and predictability as any other JAX operation. There is no hidden magic to go wrong. But it does mean the mental model is different, and that difference is worth understanding before writing any code.\nOrbax is designed around this reality. It does not try to serialize an opaque model object. It serializes pytrees, the same fundamental data structure we have been working with since\nweek 1\n.\nFrom Flax NNX object to Pytree: nnx.state\nWe spent\nweek 3\nlearning that Flax NNX gives us PyTorch style, object oriented models while secretly managing pure JAX state underneath. Checkpointing is where that \"secretly\" becomes explicit and useful.\nThe bridge function is\nnnx.state()\n. Given an NNX module, it extracts every\nnnx.Variable\ninside it, including all\nnnx.Param\nweights, into a single pytree called an\nnnx.State\n. This is exactly the kind of structure Orbax knows how to save.\nfrom flax import nnx\n\nmodel = CNN(rngs=nnx.Rngs(0))\n\n# Extract the model's state as a pytree\nstate = nnx.state(model)\nThere is a closely related pair of functions worth knowing,\nnnx.split\nand\nnnx.merge\n, which we previewed back in\nweek 3\n:\n# Split separates structure (GraphDef) from data (State)\ngraphdef, state = nnx.split(model)\n\n# Merge reconstructs a live module from structure and data\nreconstructed_model = nnx.merge(graphdef, state)\nThe\nGraphDef\ncaptures the static architecture, the layout of layers, and it never changes during training. The\nState\ncaptures the actual numbers, the weights that update every step. Orbax only ever needs to save the\nState\n. The\nGraphDef\ncan always be recreated simply by instantiating your model class again in code.\nThis distinction, static architecture versus dynamic state, is the conceptual key to everything that follows.\nThe core Orbax components\nOrbax is typically imported under the alias\nocp\n:\nimport orbax.checkpoint as ocp\nThere are a few pieces worth knowing by name.\nocp.CheckpointManager\nis the component you will use in almost every real project. It manages an entire directory of checkpoints across many training steps: deciding when to save, how many old checkpoints to retain, and how to find the most recent one when you want to resume.\nmngr = ocp.CheckpointManager(\n    directory=\"/path/to/checkpoints\",\n    options=ocp.CheckpointManagerOptions(\n        max_to_keep=3,\n        save_interval_steps=100,\n    ),\n)\nocp.Checkpointer\n, and specifically\nStandardCheckpointer\n, is the lower level component that actually handles serialization of a single pytree. In most workflows you will not touch this directly.\nCheckpointManager\nuses it internally.\nocp.args\nis a namespace of argument objects that tell Orbax exactly what kind of save or restore operation you are performing.\nocp.args.StandardSave(pytree)\nsaves a standard pytree.\nocp.args.StandardRestore(abstract_pytree)\nrestores a pytree, and requires you to supply an abstract structure describing what shape and dtype you expect back.\nocp.args.Composite(**kwargs)\nlets you save or restore several distinctly named items together, for example model parameters and optimizer state as two separate named entries in a single checkpoint.\nWe will use all three shortly.\nSaving model state: A first example\nLet's start with the simplest possible case: saving just the model's parameters at a given training step.\nimport orbax.checkpoint as ocp\nfrom flax import nnx\n\n# Assume model is an initialized nnx.Module\nmngr = ocp.CheckpointManager(\"/tmp/my_checkpoints\")\n\n# Extract the state\nstate_to_save = nnx.state(model)\n\n# Save at step 100\nmngr.save(100, args=ocp.args.StandardSave(state_to_save))\n\n# IMPORTANT: wait for the (possibly asynchronous) save to actually finish\nmngr.wait_until_finished()\n\nmngr.close()\nTwo details here deserve emphasis.\nFirst,\nmngr.wait_until_finished()\n. Orbax can save asynchronously in the background, which means your training loop does not have to stall while gigabytes of weights are written to disk. This is a meaningful performance win for large models. But it also means that if your program exits immediately after calling\nsave()\n, the write might not have actually completed yet. Always call\nwait_until_finished()\nbefore your program ends, or before you rely on that checkpoint being present.\nSecond,\nmngr.close()\n. Like any resource that manages file handles and background threads, close it cleanly when you are done, particularly at the very end of a training script.\nRestoring model state\nRestoring is slightly more involved than saving, because Orbax needs to know the exact shape and dtype structure it is restoring into before it reads a single byte from disk. This is where\nnnx.eval_shape\nbecomes essential.\nnnx.eval_shape\nbuilds an \"abstract\" version of your model: an object with the exact same structure as a real model, but where every array is replaced by a lightweight\nShapeDtypeStruct\ndescribing its shape and dtype, with no actual memory allocated. This gives Orbax a blueprint to restore into, without you needing to fully instantiate a real model first, which matters enormously for very large models where a wasted allocation could mean an out of memory error before you have even loaded anything.\n# Build an abstract model, no real memory allocated\nabstract_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))\n\n# Split into structure and abstract state\ngraphdef, abstract_state = nnx.split(abstract_model)\n\n# Ask the manager for the latest available step\nstep_to_restore = mngr.latest_step()\n\nif step_to_restore is not None:\n    restored_state = mngr.restore(\n        step_to_restore,\n        args=ocp.args.StandardRestore(abstract_state),\n    )\n    \n    # 4. Reconstruct a live, usable model\n    restored_model = nnx.merge(graphdef, restored_state)\n\nmngr.close()\nmngr.latest_step()\nis a small but genuinely useful convenience. Rather than hardcoding which step to restore, it automatically finds the most recent checkpoint written to the directory, which is exactly what you want in the common case of \"resume from wherever I last left off.\"\nThere is also a second restoration pattern worth knowing: updating an existing model object in place, rather than creating a brand new one.\nexisting_model = CNN(rngs=nnx.Rngs(0))\nnnx.update(existing_model, restored_state)\nUse\nnnx.merge\nwhen you are building a fresh model from nothing. Use\nnnx.update\nwhen you already have a live model object and simply want to overwrite its weights with restored values, for example loading a pretrained checkpoint into a model you have already partially configured.\nSaving more than just the model\nA trained model alone is not enough to resume training seamlessly. If you restore only the weights and restart with a fresh optimizer, you lose Adam's momentum and variance estimates, effectively resetting your optimizer's memory of the training trajectory so far. You also lose your exact step count, which matters if you are using a learning rate schedule tied to step number, as we built in Week 5.\nFor a real resumable training run, you want to checkpoint model parameters and optimizer state together. This is what\nocp.args.Composite\nis for.\nSaving composite state\nimport optax\nfrom flax import nnx\n\ntx = optax.adamw(learning_rate=0.001)\noptimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)\n\n# Extract each piece of state separately\nparams_state = nnx.split(optimizer, nnx.Param)[1]\noptimizer_state = nnx.state(optimizer)\n\nsave_items = {\n    \"params\": ocp.args.StandardSave(params_state),\n    \"optimizer\": ocp.args.StandardSave(optimizer_state),\n}\n\nmngr.save(\n    optimizer.step.value,\n    args=ocp.args.Composite(**save_items),\n)\nmngr.wait_until_finished()\nNotice that we key the step number off\noptimizer.step.value\n. NNX optimizers track their own step count internally, which is a convenient single source of truth for \"how far along is this training run.\"\nRestoring composite state\nRestoration follows the same abstract shape pattern as before, but now for two named items instead of one:\n# Build abstract model and optimizer\nabs_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))\nabs_optimizer = nnx.eval_shape(\n    lambda: nnx.Optimizer(abs_model, optax.adamw(0.001), wrt=nnx.Param)\n)\n\ngraphdef, abs_params_state = nnx.split(abs_model, nnx.Param)\nabs_optimizer_state = nnx.state(abs_optimizer)\n\n# Define what we're restoring into\nrestore_targets = {\n    \"params\": ocp.args.StandardRestore(abs_params_state),\n    \"optimizer\": ocp.args.StandardRestore(abs_optimizer_state),\n}\n\n# Restore\nstep = mngr.latest_step()\nrestored = mngr.restore(step, args=ocp.args.Composite(**restore_targets))\n\n# Build real, live instances and update them with restored data\nmodel_instance = CNN(rngs=nnx.Rngs(1))\noptimizer_instance = nnx.Optimizer(model_instance, optax.adamw(0.001), wrt=nnx.Param)\n\nnnx.update(model_instance, restored[\"params\"])\nnnx.update(optimizer_instance, restored[\"optimizer\"])\n\nmngr.close()\nNotice something subtle here: we initialize\nmodel_instance\nwith a fresh\nnnx.Rngs(1)\n, a completely different seed than we originally used. This is intentional and it demonstrates an important point. The random seed only matters for the initial values before restoration. Once\nnnx.update\noverwrites those values with the checkpointed weights, the original seed used to create the model is irrelevant. What matters is that the abstract structure, the shapes and dtypes, matches exactly.\nA complete, resumable training loop\nLet's bring this together into something you would actually deploy: a training loop that checks for an existing checkpoint on startup, resumes if one exists, and saves periodically as it trains.\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\nimport optax\nimport orbax.checkpoint as ocp\n\nCHECKPOINT_DIR = \"/tmp/resilient_training\"\nSAVE_EVERY = 100\nNUM_STEPS = 2000\n\ndef build_model_and_optimizer(seed=0):\n    model = CNN(rngs=nnx.Rngs(seed))\n    tx = optax.adamw(learning_rate=0.001)\n    optimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)\n    return model, optimizer\n\n@nnx.jit\ndef train_step(model, optimizer, 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    return loss\n\ndef main():\n    mngr = ocp.CheckpointManager(\n        CHECKPOINT_DIR,\n        options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=SAVE_EVERY),\n    )\n\n    model, optimizer = build_model_and_optimizer(seed=0)\n    start_step = 0\n\n    # Attempt to resume\n    latest_step = mngr.latest_step()\n    if latest_step is not None:\n        print(f\"Found existing checkpoint at step {latest_step}. Resuming.\")\n\n        abs_model = nnx.eval_shape(lambda: build_model_and_optimizer(seed=0)[0])\n        abs_optimizer = nnx.eval_shape(\n            lambda: build_model_and_optimizer(seed=0)[1]\n        )\n        graphdef, abs_params_state = nnx.split(abs_model, nnx.Param)\n        abs_optimizer_state = nnx.state(abs_optimizer)\n\n        restored = mngr.restore(\n            latest_step,\n            args=ocp.args.Composite(\n                params=ocp.args.StandardRestore(abs_params_state),\n                optimizer=ocp.args.StandardRestore(abs_optimizer_state),\n            ),\n        )\n\n        nnx.update(model, restored[\"params\"])\n        nnx.update(optimizer, restored[\"optimizer\"])\n        start_step = latest_step + 1\n    else:\n        print(\"No checkpoint found. Starting fresh.\")\n\n    # Training loop\n    for step in range(start_step, NUM_STEPS):\n        # In a real project this batch would come from the Grain\n        # pipeline built in Week 8\n        dummy_batch = {\n            \"image\": jnp.ones((32, 28, 28, 1)),\n            \"label\": jnp.zeros((32,), dtype=jnp.int32),\n        }\n\n        loss = train_step(model, optimizer, dummy_batch)\n\n        if step % SAVE_EVERY == 0:\n            params_state = nnx.split(optimizer, nnx.Param)[1]\n            optimizer_state = nnx.state(optimizer)\n\n            mngr.save(\n                step,\n                args=ocp.args.Composite(\n                    params=ocp.args.StandardSave(params_state),\n                    optimizer=ocp.args.StandardSave(optimizer_state),\n                ),\n            )\n            print(f\"Step {step:4d} | Loss: {loss:.4f} | Checkpoint saved\")\n\n    mngr.wait_until_finished()\n    mngr.close()\n    print(\"Training complete.\")\n\nif __name__ == \"__main__\":\n    main()\nRun this once, kill it midway through with\nCtrl+C\n, and run it again. You will see it detect the existing checkpoint and resume from exactly where it left off, rather than starting over. This single pattern is what separates a fragile prototype from a production ready training pipeline.\nSharded checkpointing for distributed training\nEverything above works cleanly for a model that fits comfortably on a single device. But we are building toward week 10, where we will shard models across multiple GPUs or TPUs. Checkpointing sharded state introduces one additional wrinkle worth understanding now, so it is not a surprise later.\nWhen saving, there is genuinely nothing special to do. If your state is already sharded, for example because it was produced inside a\njax.jit\ncall executed within a\nMesh\ncontext, Orbax transparently handles writing each device's shard to the correct location on disk.\nmngr.save(step, args=ocp.args.StandardSave(sharded_state_pytree))\nRestoring sharded state is where care is required. The abstract target you provide to\nStandardRestore\nmust include the sharding information describing how the restored arrays should be distributed across devices, not just their shape and dtype.\ndef create_abstract_sharded_target():\n    abstract_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))\n    _, abstract_state = nnx.split(abstract_model)\n    \n    sharding_specs = nnx.get_partition_spec(abstract_state)\n    \n    return jax.lax.with_sharding_constraint(abstract_state, sharding_specs)\n\nwith mesh:  # a jax.sharding.Mesh, covered fully in Week 10\n    abstract_target = jax.jit(create_abstract_sharded_target)()\n\nrestored_sharded_state = mngr.restore(\n    step,\n    args=ocp.args.StandardRestore(abstract_target),\n)\nThis matters because when you restore a large sharded model, you do not want Orbax to load the entire thing onto a single device and then redistribute it, which could easily exceed that single device's memory. By providing sharding information up front, Orbax can restore each shard directly to the correct device, keeping memory usage proportional to what a single device actually needs to hold.\nIt is also worth flagging a real world migration detail. Newer versions of Flax NNX changed how checkpoints structure RNG state, particularly for layers like Dropout or BatchNorm that carry their own random state. If you are restoring a checkpoint saved with an older NNX version into a newer one, this can require a specific migration step. Keep this in mind if you are ever debugging a restore that fails only for models containing stateful, randomness dependent layers.\nOther Orbax features worth knowing\nAsynchronous checkpointing\nis enabled by default in\nCheckpointManager\nand is what allows\nmngr.save()\nto return almost immediately while the actual write happens on a background thread. This is why\nwait_until_finished()\nexists and why it matters.\nAtomicity\nis guaranteed by\nCheckpointManager\n. A checkpoint write either completes fully or is treated as if it never happened. You will never end up restoring a half written, corrupted checkpoint, which is exactly the kind of failure mode you most want to avoid after a crash.\nThe TensorStore backend\nis what Orbax often uses under the hood for efficient array serialization, particularly valuable for very large arrays and for cloud storage backends where efficient, chunked I/O matters.\nCommon pitfalls\nA few mistakes are common enough to name directly.\nForgetting\nwait_until_finished()\nbefore the program exits.\nBecause saves can be asynchronous, a script that calls\nsave()\nand then immediately exits can leave you with an incomplete or missing checkpoint.\nRestoring with the wrong abstract structure.\nIf your\nabstract_state\ndoes not exactly match the structure of what you saved, for example because you changed your model's architecture between the save and the restore, Orbax will raise a clear error rather than silently loading garbage. Treat that error as useful information, not an obstacle to work around.\nNot setting\nmax_to_keep\n.\nWithout a limit,\nCheckpointManager\nwill happily keep every single checkpoint you ever save, which for a large model can consume enormous amounts of disk space over a long training run. Set\nmax_to_keep\ndeliberately based on how much history you actually need.\nSaving too frequently.\nCheckpointing has real cost, both in I/O and in the memory required to stage the write. Saving every single step is rarely necessary. A cadence tied to a meaningful interval, such as every hundred or every thousand steps, is usually the right balance.\nExercises\nSimulate a crash.\nTake the resumable training loop above, let it run for a few hundred steps, then interrupt it. Restart it and confirm the loss picks up from a similar value rather than the high loss of a freshly initialized model.\nInspect a checkpoint on disk.\nAfter saving a checkpoint, use\nview\nor your terminal to look at the directory structure Orbax creates. Note how it organizes different steps and named items.\nAdd Grain iterator state.\nUsing what you learned in Week 8, extend the training loop to also checkpoint the Grain data iterator's position, so that resuming a training run does not repeat or skip any batches.\nTest a shape mismatch.\nDeliberately change your model's\nhidden_dim\nafter saving a checkpoint, then try to restore into the new architecture. Read the resulting Orbax error carefully and confirm you understand exactly what it is telling you.\nQuick reference\nimport orbax.checkpoint as ocp\nfrom flax import nnx\n\n# Manager Setup\nmngr = ocp.CheckpointManager(\n    directory,\n    options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=100),\n)\n\n# Saving a Single Pytree\nstate = nnx.state(model)\nmngr.save(step, args=ocp.args.StandardSave(state))\nmngr.wait_until_finished()\n\n# Restoring a Single Pytree\nabstract_model = nnx.eval_shape(lambda: MyModel(rngs=nnx.Rngs(0)))\ngraphdef, abstract_state = nnx.split(abstract_model)\nrestored_state = mngr.restore(\n    mngr.latest_step(), args=ocp.args.StandardRestore(abstract_state)\n)\nrestored_model = nnx.merge(graphdef, restored_state)\n\n# Composite Save/Restore (model + optimizer)\nmngr.save(step, args=ocp.args.Composite(\n    params=ocp.args.StandardSave(params_state),\n    optimizer=ocp.args.StandardSave(optimizer_state),\n))\n\nrestored = mngr.restore(step, args=ocp.args.Composite(\n    params=ocp.args.StandardRestore(abs_params_state),\n    optimizer=ocp.args.StandardRestore(abs_optimizer_state),\n))\n\n# Updating Live Objects\nnnx.update(model, restored[\"params\"])\nnnx.update(optimizer, restored[\"optimizer\"])\n\n# Cleanup\nmngr.close()\nWhat's next\nWe can now train models that survive crashes, power outages, and preempted cloud instances. But everything we have built so far assumes a single device. Real large scale models, the kind used for foundational speech to text systems or large language models, do not fit on one GPU at all.\nNext week we tackle\nsharding and parallelism\n. We will learn how to split a model across multiple devices using\njax.sharding.Mesh\nand\nPartitionSpec\n, how to combine this with the checkpointing techniques from today, and how to think about data parallelism, tensor parallelism, and fully sharded data parallelism as complementary strategies rather than competing ones.\nNext week\n:\nSplitting the Load: Sharding Across Devices\n#\njax\n#\norbax\n#\ncheckpointing\n#\nai",
  "metadata": {
    "title": "Never lose your progress: Checkpointing with Orbax",
    "description": "Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes ou",
    "keywords": null,
    "author": null,
    "og:title": "Never lose your progress: Checkpointing with Orbax",
    "og:description": "Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes ou",
    "og:url": "https://kambale.dev/orbax",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/789d7c1c-8ad7-412c-8d7a-5cd109d97676.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": "Never lose your progress: Checkpointing with Orbax",
    "twitter:description": "Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes ou",
    "twitter:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/789d7c1c-8ad7-412c-8d7a-5cd109d97676.png",
    "article:published_time": "2026-08-07T14:40:11.974Z",
    "article:modified_time": "2026-08-07T14:40:11.974Z",
    "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/ab459134afd4/code/body-000.txt",
      "chars": 48,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-001.txt",
      "chars": 118,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-002.txt",
      "chars": 201,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-003.txt",
      "chars": 31,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-004.txt",
      "chars": 175,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-005.txt",
      "chars": 397,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-006.txt",
      "chars": 569,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-007.txt",
      "chars": 82,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-008.txt",
      "chars": 495,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-009.txt",
      "chars": 907,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-010.txt",
      "chars": 2939,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-011.txt",
      "chars": 66,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-012.txt",
      "chars": 530,
      "language": "text"
    },
    {
      "file": "attachments/ab459134afd4/code/body-013.txt",
      "chars": 1131,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/789d7c1c-8ad7-412c-8d7a-5cd109d97676.png",
      "bytes": 72886,
      "reused": true,
      "local_path": "attachments/ab459134afd4/images/000.png",
      "src": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/789d7c1c-8ad7-412c-8d7a-5cd109d97676.png",
      "alt": "Never lose your progress: Checkpointing with Orbax",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/orbax.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/why-jax-the-numpy-you-know-but-faster",
      "label": "week 1",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/flax-nnx",
      "label": "week 3",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/orbax",
      "label": "# orbax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/checkpointing",
      "label": "# checkpointing",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/ai",
      "label": "# ai",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:46Z",
  "source_url": "https://kambale.dev/orbax",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 23033,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/ab459134afd4/pages/source.html",
  "local_html": "attachments/ab459134afd4/pages/article.html",
  "body_base_path": "attachments/ab459134afd4/pages/",
  "body_markdown_path": "attachments/ab459134afd4/pages/article.md"
}