{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-1 fw=all -->\n",
    "\n",
    "# Modules and Model Construction\n",
    "\n",
    "The networks we have trained so far were small enough to write down layer by\n",
    "layer. The networks ahead are not: deep ResNets chain more than a hundred\n",
    "convolutional layers [@He.Zhang.Ren.ea.2016], and a GPT-style language\n",
    "model stacks dozens of identical Transformer blocks\n",
    "[@Radford.Wu.Child.ea.2019]. Such models are assembled from repeated\n",
    "*blocks*, groups of layers treated as units of design. The abstraction that\n",
    "makes blocks composable is the *module*.\n",
    "\n",
    "A module is an object with three responsibilities: it owns *parameters*, it\n",
    "owns *child modules*, and it implements a *forward computation* that maps\n",
    "inputs to outputs. The definition is deliberately recursive. A fully connected\n",
    "layer is a module (parameters, no children). A residual block is a module (no\n",
    "direct parameters, a few child layers that own parameters). A hundred-layer\n",
    "network is a module whose children are blocks whose children are layers. Most\n",
    "models therefore have a tree-shaped module hierarchy, as sketched in\n",
    "the figure. Reusing one child at several sites turns that tree into\n",
    "an object graph, a case we meet when tying parameters in\n",
    "that section. Almost everything\n",
    "this chapter does to a model, listing its parameters\n",
    "(that section), moving it to a GPU (that section),\n",
    "saving it to disk (that section), is implemented as a walk over\n",
    "that tree.\n",
    "\n",
    "![Layers compose into blocks and blocks compose into models, giving the usual tree-shaped module hierarchy.](../img/bg-module-tree.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-modules-and-model-construction"
   },
   "source": [
    "from dataclasses import dataclass\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-2 fw=all -->\n",
    "\n",
    "## The Module Abstraction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-3 fw=jax -->\n",
    "\n",
    "In JAX the module class is Flax's `nnx.Module`. We have used one of its\n",
    "subclasses all along: `nnx.Sequential` builds a model from a chain of layers,\n",
    "here the familiar MLP with a 256-unit ReLU hidden layer and a 10-unit output\n",
    "layer. NNX modules own their initialized parameters, so construction requires\n",
    "the layer widths and an RNG stream; a forward pass then uses the same direct\n",
    "`net(X)` notation as the other frameworks."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-1"
   },
   "source": [
    "net = nnx.Sequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,\n",
    "                     nnx.Linear(256, 10, rngs=nnx.Rngs(1)))\n",
    "\n",
    "X = jax.random.uniform(d2l.get_key(), (2, 20))\n",
    "net(X).shape"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-4 fw=jax -->\n",
    "\n",
    "`Sequential` is not a special construct. It is itself a module whose forward\n",
    "computation runs its children in order, and the children are the three modules\n",
    "we passed to it. Each child owns its parameters; selecting `nnx.Param` state\n",
    "gives a flat path view whose structure mirrors the module tree:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-2"
   },
   "source": [
    "[(path, value.shape)\n",
    " for path, value in nnx.state(net, nnx.Param).flat_state()]"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('layers', 0, 'bias'), (256,)),\n (('layers', 0, 'kernel'), (20, 256)),\n (('layers', 2, 'bias'), (10,)),\n (('layers', 2, 'kernel'), (256, 10))]"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-5 fw=jax -->\n",
    "\n",
    "`nnx.state(net, nnx.Param)` selects the trainable leaves of the object graph.\n",
    "Optimizer updates, device movement, and serialization traverse this state.\n",
    "The parameters remain owned by their layers; filters give us a state view\n",
    "without maintaining a second model representation by hand."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-3"
   },
   "source": [
    "class MLP(nnx.Module):\n",
    "    def __init__(self, rngs):\n",
    "        self.hidden = nnx.Linear(20, 256, rngs=rngs)\n",
    "        self.out = nnx.Linear(256, 10, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return self.out(nnx.relu(self.hidden(X)))"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-the-module-abstraction-4"
   },
   "source": [
    "net = MLP(nnx.Rngs(0))\n",
    "net(X).shape"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-6 fw=jax -->\n",
    "\n",
    "Two details do the work here. First, assigning `nnx.Linear` children to\n",
    "`self.hidden` and `self.out` makes them part of the NNX object graph, and their\n",
    "attribute names become paths in the selected parameter state. Second, we never\n",
    "wrote a backward method; `nnx.grad` derives gradients from whatever `__call__`\n",
    "computes. Layers are constructed once in `__init__`, keeping parameter\n",
    "creation separate from the forward computation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-7 fw=all -->\n",
    "\n",
    "## Sequential and Friends: Containers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-8 fw=jax -->\n",
    "\n",
    "To see that there is no magic left in `nnx.Sequential`, we can write it\n",
    "ourselves. Two ingredients suffice: declare a field that holds the list of\n",
    "children, and loop over them in `__call__`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-1"
   },
   "source": [
    "class MySequential(nnx.Module):\n",
    "    def __init__(self, *modules):\n",
    "        self.modules = nnx.List(modules)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        for module in self.modules:\n",
    "            X = module(X)\n",
    "        return X"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-9 fw=jax -->\n",
    "\n",
    "`nnx.List` is a graph-aware container. Both `Linear` children are tracked\n",
    "(`nnx.relu` is a plain function, with nothing to track), and their parameters\n",
    "appear beneath the `modules` path. Our version is a drop-in replacement:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-2"
   },
   "source": [
    "net = MySequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,\n",
    "                   nnx.Linear(256, 10, rngs=nnx.Rngs(1)))\n",
    "net(X).shape"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-10 fw=jax -->\n",
    "\n",
    "Graph-aware containers make registration explicit in NNX. `nnx.List` tracks\n",
    "the child modules and still supports ordinary indexing and iteration:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-3"
   },
   "source": [
    "class ListMLP(nnx.Module):\n",
    "    def __init__(self, rngs):\n",
    "        self.layers = nnx.List([nnx.Linear(20, 256, rngs=rngs),\n",
    "                                nnx.Linear(256, 10, rngs=rngs)])\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return self.layers[1](nnx.relu(self.layers[0](X)))\n",
    "\n",
    "net = ListMLP(nnx.Rngs(0))\n",
    "state = nnx.state(net, nnx.Param)\n",
    "net(X).shape, sum(x.size for _, x in state.flat_state())"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "((2, 10), 7946)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-11 fw=jax -->\n",
    "\n",
    "All 7946 parameters are present. NNX deliberately rejects a plain Python list\n",
    "that contains graph nodes, because treating such a container as static would\n",
    "silently hide its parameters. The error points directly to `nnx.List` (or\n",
    "`nnx.data`) as the repair:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-sequential-and-friends-containers-4"
   },
   "source": [
    "class PlainListSequential(nnx.Module):\n",
    "    def __init__(self, *modules):\n",
    "        self.modules = list(modules)  # Graph nodes cannot be static data\n",
    "\n",
    "try:\n",
    "    net = PlainListSequential(nnx.Linear(20, 10, rngs=nnx.Rngs(0)))\n",
    "except ValueError as e:\n",
    "    print(e)"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Found data on value of type '<class 'list'>' assigned to static attribute 'modules' of Pytree type '<class '__main__.PlainListSequential'>'. Static attributes should not contain data values. Consider one of the following options:\n\n1. If the attribute is meant to be static, remove the data values of type(s):\n\n  Linear\n\n2. If the attribute is meant to be data, wrap the value with nnx.data  on assignment:\n\n  _.modules = nnx.data(...)\n\n3. Alternatively, annotate the class attribute with nnx.data:\n\n  class PlainListSequential(Module):\n    modules: list = nnx.data()\n\n4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.\n\n5. Disable pytree for this class:\n\n  class PlainListSequential(Module, pytree=False):\n\n\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-12 fw=jax -->\n",
    "\n",
    "NNX supplies `nnx.List` and `nnx.Dict` for graph nodes. A misdeclared child\n",
    "fails at construction time instead of yielding a model that runs but trains\n",
    "nothing. Transformer implementations conventionally keep their stack of\n",
    "blocks in an `nnx.List`, with the embedding and output head as named\n",
    "attributes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-13 fw=all -->\n",
    "\n",
    "## Forward Is Just Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-14 fw=jax -->\n",
    "\n",
    "`__call__` is an ordinary Python method. Nothing restricts it to chaining\n",
    "children: it can branch, loop, call any `jnp` function, and combine\n",
    "intermediate results however it likes. The loop in `MySequential` already used\n",
    "this freedom. Its most consequential one-line use is the *residual\n",
    "connection*, the wiring idiom at the heart of ResNets and Transformers alike.\n",
    "The residual body is constructed once and then reused on every call."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-1"
   },
   "source": [
    "class ResidualBlock(nnx.Module):\n",
    "    def __init__(self, num_hiddens, rngs):\n",
    "        self.body = nnx.Sequential(\n",
    "            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs), nnx.relu,\n",
    "            nnx.Linear(num_hiddens, num_hiddens, rngs=rngs))\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return X + self.body(X)"
   ],
   "execution_count": 10,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-15 fw=all -->\n",
    "\n",
    "![The residual wiring `X + body(X)`: the input splits at a branch point into the body stack and an identity skip, and the two rejoin by addition before the block's output.](../img/bg-residual-block.svg)\n",
    "\n",
    "`X + body(X)` is not a layer any framework provides. It is plain arithmetic\n",
    "in the forward computation, and it changes what the block *is*: the block\n",
    "computes a perturbation of the identity function rather than an arbitrary\n",
    "transformation. If its body is $F$, its Jacobian is $I + J_F$, so\n",
    "backpropagation receives an additive identity contribution along the skip\n",
    "path. This contribution can still cancel against $J_F$; it is a direct path,\n",
    "not a guarantee that every gradient is preserved\n",
    "[@He.Zhang.Ren.ea.2016]. the figure diagrams\n",
    "exactly this wiring, and that section develops both points when we\n",
    "build ResNet; for now we only need the mechanics. One mechanical consequence\n",
    "is visible already: the addition forces the input and output shapes to\n",
    "agree, so a residual block has a single width that is part of its identity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-16 fw=jax -->\n",
    "\n",
    "That is why `num_hiddens` is an explicit constructor argument rather than a\n",
    "width inferred from a sample batch."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-2"
   },
   "source": [
    "block = ResidualBlock(24, nnx.Rngs(0))\n",
    "X24 = jax.random.normal(d2l.get_key(), (2, 24))\n",
    "block(X24).shape"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 24)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-17 fw=jax -->\n",
    "\n",
    "`__call__` may also use state that is neither an input nor a parameter.\n",
    "Suppose we want to damp each block's contribution by a fixed factor:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-3"
   },
   "source": [
    "class ScaledResidual(ResidualBlock):\n",
    "    def __init__(self, num_hiddens, rngs, alpha=0.5):\n",
    "        super().__init__(num_hiddens, rngs)\n",
    "        self.alpha = alpha  # Fixed architecture data, never trained\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return X + self.alpha * self.body(X)\n",
    "\n",
    "block = ScaledResidual(24, nnx.Rngs(0))\n",
    "block.alpha, [path for path, _ in nnx.state(block, nnx.Param).flat_state()]"
   ],
   "execution_count": 12,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(0.5,\n [('body', 'layers', 0, 'bias'),\n  ('body', 'layers', 0, 'kernel'),\n  ('body', 'layers', 2, 'bias'),\n  ('body', 'layers', 2, 'kernel')])"
     },
     "metadata": {},
     "execution_count": 12
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-18 fw=jax -->\n",
    "\n",
    "`alpha` enters the computation, but it is ordinary architecture data rather\n",
    "than an `nnx.Param`, so the optimizer never touches it. Constructing\n",
    "`ScaledResidual(24, nnx.Rngs(0), alpha=0.5)` reproduces it. Mutable non-parameter state,\n",
    "such as running statistics, instead uses another `nnx.Variable` subclass,\n",
    "introduced in that section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-forward-is-just-python-4"
   },
   "source": [
    "class HalvingMLP(nnx.Module):\n",
    "    def __init__(self, rngs):\n",
    "        self.proj = nnx.Linear(20, 24, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        X = self.proj(X)\n",
    "        while jnp.abs(X).sum() > 1:  # Ordinary Python control flow\n",
    "            X = X / 2\n",
    "        return X.sum()\n",
    "\n",
    "net = HalvingMLP(nnx.Rngs(0))\n",
    "net(X)"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array(-0.05208473, dtype=float32)"
     },
     "metadata": {},
     "execution_count": 13
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-19 fw=jax -->\n",
    "\n",
    "Once a model is wrapped in `jax.jit` for speed, such data-dependent Python\n",
    "control flow must be expressed with `jax.lax` primitives instead; that\n",
    "constraint arrives only with the compiler, not with the module abstraction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-20 fw=all -->\n",
    "\n",
    "## Lazy Initialization: Shapes from Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-21 fw=jax -->\n",
    "\n",
    "NNX makes a different choice: `nnx.Linear` takes both its input and output\n",
    "widths and creates its parameters in the constructor. This keeps the model\n",
    "object complete: an optimizer or checkpoint can inspect every parameter\n",
    "immediately, without a dummy forward pass. The widths also make a bad\n",
    "connection easier to diagnose, although an ordinary builder does not validate\n",
    "adjacent layers; an incompatible pair still fails when the model is called."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-22 fw=jax -->\n",
    "\n",
    "There is no uninitialized form to inspect. We state the two boundary widths\n",
    "and the hidden width once; the resulting state is available immediately:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-2"
   },
   "source": [
    "net = nnx.Sequential(nnx.Linear(20, 256, rngs=nnx.Rngs(0)), nnx.relu,\n",
    "                     nnx.Linear(256, 10, rngs=nnx.Rngs(1)))\n",
    "[(path, value.shape)\n",
    " for path, value in nnx.state(net, nnx.Param).flat_state()]"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('layers', 0, 'bias'), (256,)),\n (('layers', 0, 'kernel'), (20, 256)),\n (('layers', 2, 'bias'), (10,)),\n (('layers', 2, 'kernel'), (256, 10))]"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-23 fw=all -->\n",
    "\n",
    "Lazy layers can remove shape arithmetic from model definitions. In a\n",
    "convolutional network the flattened feature-map size depends on the input\n",
    "resolution and upstream strides. We usually avoid that fragile arithmetic by\n",
    "using global or adaptive pooling before the first linear layer. When a width\n",
    "is part of the architecture, we name it in the model config and pass it to\n",
    "adjacent layers — code that stays explicit and compact without hiding\n",
    "parameter creation behind a sample batch, whether or not the framework\n",
    "offers a lazy mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-24 fw=jax -->\n",
    "\n",
    "Parameters are available as soon as construction returns. Initialization\n",
    "randomness comes from the `nnx.Rngs` object passed to the constructors, so a\n",
    "layer's random stream is explicit (that section returns to seeding)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-25 fw=jax -->\n",
    "\n",
    "Each `nnx.Linear` accepts a `kernel_init` function and invokes it in the\n",
    "constructor. A small demonstration:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-lazy-initialization-shapes-from-data-4"
   },
   "source": [
    "class TinyMLP(d2l.Module):\n",
    "    def __init__(self, rngs):\n",
    "        super().__init__()\n",
    "        self.net = nnx.Sequential(\n",
    "            nnx.Linear(20, 256,\n",
    "                       kernel_init=nnx.initializers.xavier_uniform(),\n",
    "                       rngs=rngs), nnx.relu,\n",
    "            nnx.Linear(256, 10, rngs=rngs))\n",
    "\n",
    "model = TinyMLP(nnx.Rngs(0))\n",
    "[(path, value.shape)\n",
    " for path, value in nnx.state(model, nnx.Param).flat_state()]"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('net', 'layers', 0, 'bias'), (256,)),\n (('net', 'layers', 0, 'kernel'), (20, 256)),\n (('net', 'layers', 2, 'bias'), (10,)),\n (('net', 'layers', 2, 'kernel'), (256, 10))]"
     },
     "metadata": {},
     "execution_count": 15
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-26 fw=jax -->\n",
    "\n",
    "The constructor created every parameter and drew the first kernel from the\n",
    "Xavier initializer. Which\n",
    "initializer to use, and why Xavier's variance rule (that section)\n",
    "is a sensible default, is the subject of that section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-27 fw=all -->\n",
    "\n",
    "## Building from a Config\n",
    "\n",
    "So far every width in this section was a literal typed into a constructor.\n",
    "Real model code does not work that way. The architecture choices that vary\n",
    "across experiments, such as depth, width, and output size, must be logged with\n",
    "results and stored with checkpoints so that a saved model can be rebuilt. The\n",
    "topology remains in code; a small configuration object records its variable\n",
    "choices:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-1"
   },
   "source": [
    "@dataclass\n",
    "class MLPConfig:\n",
    "    d_in: int = 784\n",
    "    d_hidden: int = 256\n",
    "    num_blocks: int = 4\n",
    "    d_out: int = 10\n",
    "\n",
    "class ResidualMLP(nnx.Module):\n",
    "    def __init__(self, cfg, rngs):\n",
    "        self.cfg = cfg\n",
    "        self.proj = nnx.Linear(cfg.d_in, cfg.d_hidden, rngs=rngs)\n",
    "        self.blocks = nnx.List([\n",
    "            ResidualBlock(cfg.d_hidden, rngs) for _ in range(cfg.num_blocks)])\n",
    "        self.out = nnx.Linear(cfg.d_hidden, cfg.d_out, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        X = self.proj(X)\n",
    "        for block in self.blocks:\n",
    "            X = block(X)\n",
    "        return self.out(X)\n",
    "\n",
    "def build(cfg: MLPConfig) -> nnx.Module:\n",
    "    return ResidualMLP(cfg, nnx.Rngs(0))"
   ],
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-28 fw=jax -->\n",
    "\n",
    "One config produces one architecture. The input width is explicit because\n",
    "NNX creates parameters in the constructor. `nnx.List` registers the repeated\n",
    "residual blocks, while the loop in `__call__` controls their execution.\n",
    "Selecting parameter state displays the resulting object graph:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-2"
   },
   "source": [
    "net = build(MLPConfig())\n",
    "[(path, value.shape)\n",
    " for path, value in nnx.state(net, nnx.Param).flat_state()]"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "[(('blocks', 0, 'body', 'layers', 0, 'bias'), (256,)),\n (('blocks', 0, 'body', 'layers', 0, 'kernel'), (256, 256)),\n (('blocks', 0, 'body', 'layers', 2, 'bias'), (256,)),\n (('blocks', 0, 'body', 'layers', 2, 'kernel'), (256, 256)),\n (('blocks', 1, 'body', 'layers', 0, 'bias'), (256,)),\n (('blocks', 1, 'body', 'layers', 0, 'kernel'), (256, 256)),\n (('blocks', 1, 'body', 'layers', 2, 'bias'), (256,)),\n (('blocks', 1, 'body', 'layers', 2, 'kernel'), (256, 256)),\n (('blocks', 2, 'body', 'layers', 0, 'bias'), (256,)),\n (('blocks', 2, 'body', 'layers', 0, 'kernel'), (256, 256)),\n (('blocks', 2, 'body', 'layers', 2, 'bias'), (256,)),\n (('blocks', 2, 'body', 'layers', 2, 'kernel'), (256, 256)),\n (('blocks', 3, 'body', 'layers', 0, 'bias'), (256,)),\n (('blocks', 3, 'body', 'layers', 0, 'kernel'), (256, 256)),\n (('blocks', 3, 'body', 'layers', 2, 'bias'), (256,)),\n (('blocks', 3, 'body', 'layers', 2, 'kernel'), (256, 256)),\n (('out', 'bias'), (10,)),\n (('out', 'kernel'), (256, 10)),\n (('proj', 'bias'), (256,)),\n (('proj', 'kernel'), (784, 256))]"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-3"
   },
   "source": [
    "net(jax.random.uniform(d2l.get_key(), (2, 784))).shape"
   ],
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(2, 10)"
     },
     "metadata": {},
     "execution_count": 18
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-29 fw=all -->\n",
    "\n",
    "The variable architecture choices are now data. Rescaling the model is a change to two fields, not\n",
    "an edit to model code:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "model-construction-building-from-a-config-4"
   },
   "source": [
    "for cfg in (MLPConfig(), MLPConfig(d_hidden=512, num_blocks=8)):\n",
    "    net = build(cfg)\n",
    "    n = sum(x.size for _, x in nnx.state(net, nnx.Param).flat_state())\n",
    "    print(f'd_hidden={cfg.d_hidden}, num_blocks={cfg.num_blocks}: '\n",
    "          f'{n:,} parameters')"
   ],
   "execution_count": 19,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "d_hidden=256, num_blocks=4: 729,866 parameters\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "d_hidden=512, num_blocks=8: 4,609,546 parameters\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-30 fw=jax -->\n",
    "\n",
    "Because `build` is deterministic in `cfg`, the config is all we need to\n",
    "reconstruct the object graph later; that section saves it\n",
    "alongside the state. Width and depth fields feeding a loop that stacks\n",
    "identical residual blocks form a common construction pattern in the\n",
    "Transformer models later in the book."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-31 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-32 fw=jax -->\n",
    "\n",
    "A module owns parameters, child modules, and a `__call__` method. Layers,\n",
    "blocks, and whole models are the same kind of object, and state selection,\n",
    "optimizer updates, and serialization traverse the NNX object graph.\n",
    "`nnx.List` and `nnx.Dict` register collections of children; a plain container\n",
    "holding graph nodes is rejected. `__call__` is ordinary Python, so a residual\n",
    "connection is one line. NNX linear layers take both widths and create their\n",
    "parameters in the constructor. A dataclass config plus a deterministic build\n",
    "function records those widths and repeated-block counts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-33 fw=all -->\n",
    "\n",
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=model-construction-md-34 fw=jax -->\n",
    "\n",
    "1. Extend `ListMLP` with an `nnx.Dict` of submodules and verify that every\n",
    "   parameter appears in `nnx.state(net, nnx.Param)`. Then replace it with a\n",
    "   plain dictionary holding those modules and describe how NNX reports the\n",
    "   invalid graph container.\n",
    "1. Implement a `ParallelBlock` that takes two child modules `net1` and `net2`\n",
    "   as fields, runs both on the same input, and concatenates their outputs\n",
    "   along the last dimension. What must be true of the two children's outputs\n",
    "   for the concatenation to be valid?\n",
    "1. Extend `MLPConfig` with an activation switch (for example,\n",
    "   `act: str = 'relu'`) and make `ResidualMLP` honor it. Which decisions belong\n",
    "   in the config and which belong in code? Where would you put a choice between\n",
    "   `ResidualBlock` and a plain feed-forward block?\n",
    "1. `ResidualBlock` requires its input and output widths to agree. Suppose you\n",
    "   want a block whose output is wider than its input. Give two standard fixes\n",
    "   and the cost of each. (that section uses one of them in ResNet.)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
