{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-1 fw=all -->\n",
    "\n",
    "# Saving, Loading, and Pretrained Weights\n",
    "\n",
    "A trained network is two separate things kept in two separate places. The\n",
    "*code* is the class you wrote: its layers, its `forward` pass, the config that\n",
    "sized it. The *model state* is the tensors that training filled in: weights,\n",
    "biases, and running statistics. Optimizer moments, random-number streams, and\n",
    "the step belong to the wider *training state*. A weight file saves model state;\n",
    "a resumable checkpoint saves full training state. The code\n",
    "stays in your source repository, under version control, exactly like any other\n",
    "Python. To bring a model back you need both halves: run the code to rebuild an\n",
    "empty network, then pour the saved state into it.\n",
    "\n",
    "This split explains most of what follows. It is why a checkpoint cannot\n",
    "resurrect a model on its own, why the config object from\n",
    "that section belongs *inside* the checkpoint, and why the\n",
    "format that stores the state matters once you start sharing files with people who\n",
    "do not have your code."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-saving-loading-and-pretrained-weights"
   },
   "source": [
    "import json\n",
    "import os\n",
    "import struct\n",
    "from dataclasses import asdict, dataclass\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "import optax\n",
    "import orbax.checkpoint as ocp\n",
    "from safetensors.flax import load_file, save_file\n",
    "from d2l import jax as d2l\n",
    "from d2l.nnx_resnet import ResNet50"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-2 fw=all -->\n",
    "\n",
    "## State, Not Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-3 fw=jax -->\n",
    "\n",
    "The state of a network is a tree of typed, named variables introduced in\n",
    "that section. Before we save a whole model, the warm-up is that\n",
    "`jnp.save` and `jnp.load` work on any array, and, through NumPy's pickle\n",
    "fallback, on the dicts that hold them."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-1"
   },
   "source": [
    "x = jnp.arange(4)\n",
    "jnp.save('tensors.npy', {'x': x, 'y': jnp.zeros(4)})\n",
    "jnp.load('tensors.npy', allow_pickle=True).item()"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'x': Array([0, 1, 2, 3], dtype=int32),\n 'y': Array([0., 0., 0., 0.], dtype=float32)}"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-4 fw=jax -->\n",
    "\n",
    "A model's parameter state is one such structure. Its paths follow the module\n",
    "graph (`hidden.kernel`, `output.bias`; Flax calls a weight matrix a `kernel`).\n",
    "Here is the tree for a small MLP."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-state-not-code-2"
   },
   "source": [
    "class MLP(nnx.Module):\n",
    "    def __init__(self, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.hidden = nnx.Linear(20, 256, rngs=rngs)\n",
    "        self.output = nnx.Linear(256, 10, rngs=rngs)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.output(nnx.relu(self.hidden(x)))\n",
    "\n",
    "net = MLP()\n",
    "X = jax.random.normal(d2l.get_key(), (2, 20))\n",
    "Y = net(X)\n",
    "params = nnx.state(net, nnx.Param)\n",
    "[(path, tuple(value.shape)) for path, value in params.flat_state()]"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('hidden', 'bias'), (256,)),\n (('hidden', 'kernel'), (20, 256)),\n (('output', 'bias'), (10,)),\n (('output', 'kernel'), (256, 10))]"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-5 fw=all -->\n",
    "\n",
    "Nothing in this dictionary knows it came from a class called `MLP`. The names\n",
    "and shapes refill a network with a compatible architecture and naming\n",
    "structure. Renaming an attribute or changing the nesting can break that match\n",
    "even when the computation is equivalent; the final section shows how to\n",
    "migrate such keys explicitly.\n",
    "\n",
    "## safetensors: the Interchange Format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-6 fw=jax -->\n",
    "\n",
    "`jnp.save` writes NumPy's `.npy` format. For a single array that is pure data:\n",
    "a small header with the dtype and shape, then the raw bytes. The warm-up dict is\n",
    "another matter. NumPy can only store a dict by falling back to Python's\n",
    "`pickle`, which does not store data so much as a program that *reconstructs*\n",
    "data, and loading runs that program; that is what `allow_pickle=True` opted\n",
    "into. For a file you wrote and never let out of your control this is harmless.\n",
    "For a file you downloaded it is a remote-code-execution vector: any object in\n",
    "the stream can name a callable for the loader to invoke, which is why NumPy\n",
    "refuses pickled contents by default (`allow_pickle=False`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-7 fw=jax -->\n",
    "\n",
    "`allow_pickle=False` is a refusal, not a fix: it keeps the loader safe by\n",
    "declining to load the very files, dicts of named parameters, that model sharing\n",
    "needs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-8 fw=all -->\n",
    "\n",
    "safetensors removes the problem at the root by\n",
    "having no program to run. As the figure lays out byte\n",
    "by byte, a safetensors file is an 8-byte little-endian integer giving the\n",
    "header length, then that many bytes of JSON naming each tensor's dtype, shape,\n",
    "and byte range, then the raw tensor bytes back to back. There is no opcode\n",
    "stream to interpret, so the format itself carries no executable program.\n",
    "Loaders must still validate malformed headers, offsets, and allocation sizes.\n",
    "It is also framework-neutral and memory-mappable, which is why model hubs\n",
    "commonly prefer it.\n",
    "Save and reload the MLP's state through it and confirm the round trip is exact.\n",
    "\n",
    "![The safetensors file as one horizontal byte strip: an 8-byte header length, a JSON header naming each tensor's dtype, shape, and byte offsets, and the raw tensor bytes packed back to back with no gaps, with two of the file's own data_offsets entries traced down to their exact span in the bar.](../img/bg-safetensors-layout.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-9 fw=jax -->\n",
    "\n",
    "One mismatch to bridge first: safetensors stores a flat mapping from names to\n",
    "tensors, while NNX state forms a nested pytree. This MLP has only string-keyed\n",
    "paths, so two small helpers join them with dots and split them on the way back.\n",
    "For an arbitrary pytree, preserve typed list indices and escape dots in keys,\n",
    "or restore values against a template instead of using this shortcut."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-3"
   },
   "source": [
    "def flatten(tree):\n",
    "    if hasattr(tree, 'flat_state'):\n",
    "        leaves = tree.flat_state()\n",
    "    else:\n",
    "        leaves = jax.tree_util.tree_flatten_with_path(tree)[0]\n",
    "        leaves = [(tuple(getattr(k, 'key', getattr(k, 'idx', k))\n",
    "                         for k in path), value)\n",
    "                  for path, value in leaves]\n",
    "    return {'.'.join(map(str, path)): value for path, value in leaves}\n",
    "\n",
    "def unflatten(flat):\n",
    "    tree = {}\n",
    "    for name, value in flat.items():\n",
    "        *parents, leaf = name.split('.')\n",
    "        node = tree\n",
    "        for p in parents:\n",
    "            node = node.setdefault(p, {})\n",
    "        node[leaf] = value\n",
    "    return tree\n",
    "\n",
    "save_file(flatten(params), 'mlp-jax.safetensors')\n",
    "restored = unflatten(load_file('mlp-jax.safetensors'))\n",
    "clone = MLP(nnx.Rngs(1))\n",
    "clone_state = nnx.state(clone, nnx.Param)\n",
    "nnx.replace_by_pure_dict(clone_state, restored)\n",
    "nnx.update(clone, clone_state)\n",
    "exact = jax.tree_util.tree_all(jax.tree_util.tree_map(\n",
    "    jnp.array_equal, restored, nnx.to_pure_dict(params)))\n",
    "exact, jnp.array_equal(clone(X), Y)"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(True, Array(True, dtype=bool))"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-10 fw=all -->\n",
    "\n",
    "Because the header is plain JSON at a known offset, you can read it without the\n",
    "library and see there is no magic to the format."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-safetensors-the-interchange-format-4"
   },
   "source": [
    "with open('mlp-jax.safetensors', 'rb') as f:\n",
    "    n = struct.unpack('<Q', f.read(8))[0]   # header length, little-endian\n",
    "    header = json.loads(f.read(n))\n",
    "header['hidden.kernel']"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "{'dtype': 'F32', 'shape': [20, 256], 'data_offsets': [1024, 21504]}"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-11 fw=jax -->\n",
    "\n",
    "`jnp.save` keeps its place for your own scratch arrays and quick experiments.\n",
    "safetensors is what you use to hand a model to anyone else."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-12 fw=all -->\n",
    "\n",
    "## Checkpointing a Training Run\n",
    "\n",
    "A checkpoint you can resume from holds more than weights. Resuming means picking\n",
    "up the optimizer where it stopped, and Adam's state is the running first and\n",
    "second moments of the gradients from that section. Drop them and\n",
    "the optimizer restarts its momentum from zero, so the first steps after a resume\n",
    "no longer behave like a continuation. A full checkpoint therefore bundles the\n",
    "model state, the optimizer state, the RNG state (so data shuffling and dropout\n",
    "continue the same stream), the step counter, and the config that sizes the model\n",
    "when you rebuild it. the figure pairs each of those\n",
    "five compartments with the exact thing it restores on resume.\n",
    "\n",
    "![A checkpoint file's five compartments, each paired by an arrow with what it restores on resume: model state_dict with weights, optimizer state with momentum and second moments, RNG state with data order and dropout, step with schedule position, and config with architecture.](../img/bg-checkpoint-contents.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-13 fw=jax -->\n",
    "\n",
    "NNX exposes the model and optimizer state as pytrees, including the optimizer's\n",
    "step counter. Orbax, the JAX checkpointing library, saves and restores such trees whole: its\n",
    "`StandardCheckpointer` commits a temporary directory by rename. The demonstration\n",
    "overwrites a fixed path with `force=True` so the notebook is rerunnable;\n",
    "production runs should write numbered steps through `CheckpointManager` and\n",
    "retain the previous completed step. Because these natives\n",
    "already cover the job, the jax tab defines no helper of its own; the calls below\n",
    "are the idiom as you would write it in any project. The config rides along as\n",
    "one more branch of the saved tree, as does the PRNG key, which in JAX is\n",
    "explicit data rather than hidden global state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-14 fw=all -->\n",
    "\n",
    "Train a tiny regressor for a hundred steps and checkpoint it. The `Config`\n",
    "dataclass is what a rebuild reads to size the model, so it travels with the\n",
    "weights."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-2"
   },
   "source": [
    "@dataclass\n",
    "class Config:\n",
    "    in_dim: int = 20\n",
    "    hidden: int = 64\n",
    "    lr: float = 0.05\n",
    "\n",
    "def build(cfg):\n",
    "    return nnx.Sequential(\n",
    "        nnx.Linear(cfg.in_dim, cfg.hidden, rngs=nnx.Rngs(0)), nnx.relu,\n",
    "        nnx.Linear(cfg.hidden, 1, rngs=nnx.Rngs(1)))\n",
    "\n",
    "def fresh_state(cfg):\n",
    "    model = build(cfg)\n",
    "    optimizer = nnx.Optimizer(model, optax.adam(cfg.lr), wrt=nnx.Param)\n",
    "    return model, optimizer\n",
    "\n",
    "cfg = Config()\n",
    "model = build(cfg)\n",
    "data = jax.random.normal(d2l.get_key(), (256, cfg.in_dim))\n",
    "target = (data @ jax.random.normal(d2l.get_key(), (cfg.in_dim, 1))\n",
    "          + 0.1 * jax.random.normal(d2l.get_key(), (256, 1)))\n",
    "\n",
    "def loss(model):\n",
    "    return jnp.mean((model(data) - target) ** 2)\n",
    "\n",
    "@nnx.jit\n",
    "def step(model, optimizer):\n",
    "    l, grads = nnx.value_and_grad(loss)(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return l\n",
    "\n",
    "model, optimizer = fresh_state(cfg)\n",
    "train_key = jax.random.key(1)\n",
    "for _ in range(100):\n",
    "    l = step(model, optimizer)\n",
    "\n",
    "ckptr = ocp.StandardCheckpointer()\n",
    "ckptr.save(os.path.abspath('run-jax'),        # orbax wants an absolute path\n",
    "           {'model': nnx.to_pure_dict(nnx.state(model)),\n",
    "            'optimizer': nnx.to_pure_dict(nnx.state(optimizer)),\n",
    "            'rng': train_key,\n",
    "            'cfg': asdict(cfg)}, force=True)\n",
    "int(optimizer.step), round(float(loss(model)), 4)"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(100, 0.0021)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-15 fw=all -->\n",
    "\n",
    "The restore is exact. Corrupt every parameter, load the checkpoint back, and the\n",
    "loss returns to where it was."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-16 fw=jax -->\n",
    "\n",
    "Orbax fills a template with the saved arrays. We then update freshly built NNX\n",
    "objects from the restored pure dictionaries. Rebuilding first checks that the\n",
    "architecture still agrees with the checkpoint structure."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-3"
   },
   "source": [
    "model_state = nnx.state(model)\n",
    "nnx.update(model, jax.tree.map(lambda p: p + 1.0, model_state))\n",
    "before = float(loss(model))\n",
    "template_model, template_optimizer = fresh_state(cfg)\n",
    "template = {\n",
    "    'model': nnx.to_pure_dict(nnx.state(template_model)),\n",
    "    'optimizer': nnx.to_pure_dict(nnx.state(template_optimizer)),\n",
    "    'rng': jax.random.key(0),\n",
    "    'cfg': asdict(Config())}\n",
    "ckpt = ckptr.restore(os.path.abspath('run-jax'), template)\n",
    "cfg = Config(**ckpt['cfg'])\n",
    "train_key = ckpt['rng']\n",
    "restored_model = nnx.state(model)\n",
    "nnx.replace_by_pure_dict(restored_model, ckpt['model'])\n",
    "nnx.update(model, restored_model)\n",
    "restored_optimizer = nnx.state(optimizer)\n",
    "nnx.replace_by_pure_dict(restored_optimizer, ckpt['optimizer'])\n",
    "nnx.update(optimizer, restored_optimizer)\n",
    "after = float(loss(model))\n",
    "f'perturbed {before:.2f} -> restored {after:.4f}'"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'perturbed 50460.26 -> restored 0.0021'"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-17 fw=all -->\n",
    "\n",
    "Now the reason the optimizer state is in the file. Resume the run two ways from\n",
    "the same checkpoint: once restoring the optimizer, once with a fresh one holding\n",
    "only the weights. The network is near its minimum, so the correct continuation\n",
    "barely moves. A fresh Adam, with its moment estimates reset and its bias\n",
    "correction starting over, takes an oversized first step and overshoots."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-checkpointing-a-training-run-4"
   },
   "source": [
    "full, full_opt = fresh_state(cfg)\n",
    "full_state, full_opt_state = nnx.state(full), nnx.state(full_opt)\n",
    "nnx.replace_by_pure_dict(full_state, ckpt['model'])\n",
    "nnx.replace_by_pure_dict(full_opt_state, ckpt['optimizer'])\n",
    "nnx.update(full, full_state)\n",
    "nnx.update(full_opt, full_opt_state)\n",
    "full_losses = []\n",
    "for _ in range(5):\n",
    "    l = step(full, full_opt)\n",
    "    full_losses.append(round(float(l), 4))\n",
    "\n",
    "fresh, fresh_opt = fresh_state(cfg)\n",
    "fresh_state_ = nnx.state(fresh)\n",
    "nnx.replace_by_pure_dict(fresh_state_, ckpt['model'])\n",
    "nnx.update(fresh, fresh_state_)\n",
    "fresh_losses = []\n",
    "for _ in range(5):\n",
    "    l = step(fresh, fresh_opt)\n",
    "    fresh_losses.append(round(float(l), 4))\n",
    "\n",
    "print('full  optimizer:', full_losses)\n",
    "print('fresh optimizer:', fresh_losses)"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "full  optimizer: [0.0021, 0.002, 0.002, 0.002, 0.002]\nfresh optimizer: [0.0021, 3.0551, 0.2507, 0.8011, 1.4134]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-18 fw=all -->\n",
    "\n",
    "The full-state run keeps descending; the weights-only run spikes and has to claw\n",
    "its way back. That transient is the cost of forgetting the optimizer, and it is\n",
    "why \"just the weights\" is not a resumable checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-19 fw=jax -->\n",
    "\n",
    "For models too large to hold in memory, orbax already works at scale: a\n",
    "checkpoint is a directory of per-array files rather than one monolith, a restore\n",
    "can target a device sharding so each accelerator materializes only its own\n",
    "pieces, and multi-host jobs save and restore in parallel;\n",
    "that section returns to the machinery when models get that big."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-20 fw=all -->\n",
    "\n",
    "## Loading Weights You Did Not Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-21 fw=jax -->\n",
    "\n",
    "The most common reason to load parameters is that someone else produced them.\n",
    "You take a network trained on a large dataset and adapt it: keep the learned\n",
    "feature extractor, replace the final layer for your own labels. The book's NNX\n",
    "ResNet-50 loader downloads pinned Microsoft ImageNet safetensors from the\n",
    "Hugging Face Hub and maps them into an NNX model. We use that real artifact,\n",
    "then reuse the small MLP file to show the key-by-key surgery without burying the\n",
    "mechanism under a ResNet-sized state tree."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-1"
   },
   "source": [
    "pretrained_net = ResNet50.from_pretrained()\n",
    "pretrained_net.fc = nnx.Linear(2048, 10, rngs=nnx.Rngs(2))\n",
    "print('new pretrained-model head:', pretrained_net.fc.kernel.shape)\n",
    "\n",
    "class Classifier(nnx.Module):\n",
    "    def __init__(self, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.hidden = nnx.Linear(20, 256, rngs=rngs)\n",
    "        self.head = nnx.Linear(256, 2, rngs=rngs)  # new head, new name\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.head(nnx.relu(self.hidden(x)))\n",
    "\n",
    "classifier = Classifier()\n",
    "new_params = nnx.state(classifier, nnx.Param)\n",
    "[(path, tuple(value.shape)) for path, value in new_params.flat_state()]"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "Downloading bytes:           |  0.00B            "
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "Reconstructing (incomplete total...): |          |  0.00B /  0.00B            "
     },
     "metadata": {}
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "Fetching 1 files:   0%|          | 0/1 [00:00<?, ?it/s]"
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "new pretrained-model head: (2048, 10)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('head', 'bias'), (2,)),\n (('head', 'kernel'), (256, 2)),\n (('hidden', 'bias'), (256,)),\n (('hidden', 'kernel'), (20, 256))]"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-22 fw=jax -->\n",
    "\n",
    "There is no `strict=False` to lean on: the merge is yours to write, and so is\n",
    "the report. Take from the file every entry whose name and shape match the new\n",
    "model, keep the fresh initialization for the rest, and compute the two key sets\n",
    "that say what happened: *missing*, parameters the model has but the file did not\n",
    "fill, and *unexpected*, file entries with no home in the model."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-2"
   },
   "source": [
    "file_flat = load_file('mlp-jax.safetensors')\n",
    "new_flat = flatten(new_params)\n",
    "matched = {k: v for k, v in file_flat.items()\n",
    "           if k in new_flat and v.shape == new_flat[k].shape}\n",
    "merged = unflatten({**new_flat, **matched})\n",
    "print('missing:', sorted(set(new_flat) - set(matched)))\n",
    "print('unexpected:', sorted(set(file_flat) - set(new_flat)))\n",
    "nnx.replace_by_pure_dict(new_params, merged)\n",
    "nnx.update(classifier, new_params)\n",
    "print('loaded trunk:', jnp.array_equal(\n",
    "    classifier.hidden.kernel, restored['hidden']['kernel']))"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "missing: ['head.bias', 'head.kernel']\nunexpected: ['output.bias', 'output.kernel']\nloaded trunk: True\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-23 fw=jax -->\n",
    "\n",
    "Read both sets before trusting the merged tree. The two `head` entries under\n",
    "*missing* are expected: that layer is new and meant to start random. The two\n",
    "`output` entries under *unexpected* are the file's old head, stranded by the\n",
    "rename, and a renamed layer is what *unexpected* usually means in practice. The\n",
    "lesson is the same as with a built-in report, only stricter, because nothing\n",
    "prints it unless you do: name which keys you expect in each set and treat\n",
    "anything else as a bug."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "saving-loading-loading-weights-you-did-not-train-3"
   },
   "source": [
    "def labels(params):\n",
    "    return jax.tree_util.tree_map_with_path(\n",
    "        lambda path, _: ('train' if any(\n",
    "            getattr(p, 'key', None) == 'head' for p in path) else 'freeze'),\n",
    "        params)\n",
    "\n",
    "tx = optax.multi_transform(\n",
    "    {'train': optax.adam(0.05), 'freeze': optax.set_to_zero()}, labels)\n",
    "optimizer = nnx.Optimizer(classifier, tx, wrt=nnx.Param)\n",
    "\n",
    "hidden_before = classifier.hidden.kernel[...].copy()\n",
    "head_before = classifier.head.kernel[...].copy()\n",
    "\n",
    "@nnx.jit\n",
    "def frozen_step(model, optimizer, X):\n",
    "    loss, grads = nnx.value_and_grad(\n",
    "        lambda m: jnp.square(m(X)).mean())(model)\n",
    "    optimizer.update(model, grads)\n",
    "    return loss\n",
    "\n",
    "_ = frozen_step(classifier, optimizer, X)\n",
    "print('trunk unchanged:', jnp.array_equal(\n",
    "    hidden_before, classifier.hidden.kernel))\n",
    "print('head changed:', not jnp.array_equal(head_before, classifier.head.kernel))\n",
    "\n",
    "sizes = flatten(jax.tree_util.tree_map(jnp.size, merged))\n",
    "flat_labels = flatten(labels(merged))\n",
    "trainable = sum(int(s) for k, s in sizes.items() if flat_labels[k] == 'train')\n",
    "total = sum(int(s) for s in sizes.values())\n",
    "f'{trainable} trainable of {total}'"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "trunk unchanged: True\nhead changed: True\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'514 trainable of 5890'"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-24 fw=jax -->\n",
    "\n",
    "The Hugging Face Hub distributes JAX weights as safetensors, so the flat dict\n",
    "you just merged has the same shape as the artifact you will download in\n",
    "practice. This section covers *how* to load and adapt pretrained weights;\n",
    "that section covers when it helps and how far to unfreeze."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-25 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-26 fw=jax -->\n",
    "\n",
    "A saved model is state, not code: a pytree of named arrays that means something\n",
    "only once the code that built the network runs again. For your own files\n",
    "`jnp.save` is fine; for files you share, safetensors stores the same tensors\n",
    "behind a plain JSON header, with no pickle to execute, which is why hubs\n",
    "standardize on it. A resumable checkpoint contains model state, optimizer\n",
    "state, and step in pytrees that Orbax saves atomically in a single call, or a\n",
    "resume restarts the optimizer's momentum from zero. Loading someone else's\n",
    "weights is pytree surgery, and the missing/unexpected sets you compute are a\n",
    "diagnostic to read rather than a formality to skip."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=saving-loading-md-27 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Even if you never deploy to another machine, name two reasons to checkpoint.\n",
    "   Then consider the atomic write: if the checkpoint were written straight to its\n",
    "   final path (delete the `os.replace` from `save_checkpoint`, the\n",
    "   rename-into-place that orbax performs, or the fresh numbered files that\n",
    "   `tf.train.Checkpoint.save` writes), describe the failure a crash mid-write\n",
    "   now causes, and why the atomic version avoids it.\n",
    "1. Read the first 8 bytes of the safetensors file you wrote for the MLP as a\n",
    "   little-endian integer, as the header cell does. How large is the JSON header\n",
    "   for the MLP, and how does it grow if you double the hidden width?\n",
    "1. Save the MLP's parameters cast to `bfloat16` and load them back into a\n",
    "   `float32` model (that section). What is lost? Is that acceptable\n",
    "   for inference? For resuming training?\n",
    "1. Take two checkpoints of the regressor 50 steps apart, average their weight\n",
    "   tensors into a third set of parameters, and evaluate it. The result previews\n",
    "   weight averaging."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
