{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-1 fw=all -->\n",
    "\n",
    "# Initialization\n",
    "\n",
    "Every parameter in a network has a value before training starts: a fully\n",
    "connected layer mapping 512 inputs to 256 outputs brings a $256 \\times 512$\n",
    "weight matrix filled with random numbers. that section explains why the scale of those\n",
    "numbers decides whether a deep network trains at all, and derives the Xavier\n",
    "[@Glorot.Bengio.2010] and He [@He.Zhang.Ren.ea.2015] schemes that\n",
    "keep signal variance stable across layers. This section is the practical\n",
    "companion: what the library does when you say nothing, how to impose a scheme\n",
    "of your own on a whole model or a single block, what transformer-era code adds\n",
    "on top of Xavier and He, and how to write an initializer that no menu provides."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-initialization"
   },
   "source": [
    "import math\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=init-md-2 fw=all -->\n",
    "\n",
    "## Defaults and When to Override Them\n",
    "\n",
    "You can usually ignore initialization because the default is sensible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-3 fw=jax -->\n",
    "\n",
    "NNX creates parameters when `nnx.Linear` is constructed, using the explicit\n",
    "input width to determine fan-in. Its default is LeCun normal: the weight comes\n",
    "from a normal distribution with standard deviation\n",
    "$1/\\sqrt{n_\\textrm{in}}$, truncated so that no entry lands in the far tails\n",
    "(the initializer rescales after clipping to keep the standard deviation on target, which\n",
    "puts the hard bound near $2.3$ standard deviations), and the bias starts at\n",
    "zero. Both claims are cheap to check against a fresh layer:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-defaults-and-when-to-override-them"
   },
   "source": [
    "layer = nnx.Linear(512, 256, rngs=nnx.Rngs(0))\n",
    "w = layer.kernel\n",
    "print(f'std {w.std():.4f}, predicted {1 / math.sqrt(512):.4f}')\n",
    "print(f'max |w| {jnp.abs(w).max():.4f}, '\n",
    "      f'bias all zero: {bool((layer.bias == 0).all())}')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "std 0.0443, predicted 0.0442\nmax |w| 0.1005, bias all zero: True\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-4 fw=all -->\n",
    "\n",
    "Other libraries make different but equally fan-aware choices, with one legacy\n",
    "exception:\n",
    "\n",
    "| Library | Default for a dense layer's weight |\n",
    "|:--|:--|\n",
    "| PyTorch | $U(-1/\\sqrt{n_\\textrm{in}},\\, 1/\\sqrt{n_\\textrm{in}})$, bias likewise |\n",
    "| Flax (JAX) | LeCun normal: truncated normal, std $1/\\sqrt{n_\\textrm{in}}$; zero bias |\n",
    "| Keras (TensorFlow) | Glorot uniform: $U(\\pm\\sqrt{6/(n_\\textrm{in} + n_\\textrm{out})})$; zero bias |\n",
    "| MXNet Gluon | $U(-0.07, 0.07)$, not fan-aware, so override it; zero bias |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-5 fw=jax -->\n",
    "\n",
    "When should you override? Four situations recur: a deep network without\n",
    "normalization layers, where variance compounds across depth (we demonstrate\n",
    "this below); reproducing a paper whose results depend on its initialization\n",
    "recipe; architecture-specific corrections such as the residual scaling later\n",
    "in this section; and parameters you create by hand, though here NNX forces\n",
    "the choice on you: construct the value explicitly and wrap it in `nnx.Param`.\n",
    "NNX layers receive their input and output widths in the constructor, and\n",
    "their initializers run there as well; no later `init` pass infers fan-in from\n",
    "the first input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-6 fw=all -->\n",
    "\n",
    "## Applying Initializers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-7 fw=jax -->\n",
    "\n",
    "NNX affine layers such as `Linear` and `Conv` accept `kernel_init` and\n",
    "`bias_init`, each a function\n",
    "`(key, shape, dtype) -> array`, and `nnx.initializers` supplies the standard\n",
    "menu. These functions run in the constructor. For a model that already\n",
    "exists, we can walk its modules and assign new values to selected parameters.\n",
    "For a parameter in a custom module, call an initializer directly and wrap its\n",
    "result in `nnx.Param`. We show the first two routes here."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-applying-initializers-1"
   },
   "source": [
    "init_normal = nnx.initializers.normal(0.01)\n",
    "net = nnx.Sequential(\n",
    "    nnx.Linear(4, 8, kernel_init=init_normal,\n",
    "               bias_init=nnx.initializers.zeros, rngs=nnx.Rngs(0)),\n",
    "    nnx.relu,\n",
    "    nnx.Linear(8, 1, kernel_init=init_normal,\n",
    "               bias_init=nnx.initializers.zeros, rngs=nnx.Rngs(1)))\n",
    "\n",
    "net.layers[0].kernel[:, 0], net.layers[0].bias[0]"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([0.01004014, 0.01661628, 0.00595053, 0.00908357], dtype=float32),\n Array(0., dtype=float32))"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-8 fw=jax -->\n",
    "\n",
    "Mixing schemes per block needs no dispatch at all: each layer names its own\n",
    "initializer, so the choice sits exactly where the layer is declared. Below we\n",
    "give the first layer Xavier initialization (that section) and set\n",
    "the last to a constant, a poor idea for training by the symmetry argument of\n",
    "that section, but it makes the mechanics visible:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-applying-initializers-2"
   },
   "source": [
    "net = nnx.Sequential(\n",
    "    nnx.Linear(4, 8, kernel_init=nnx.initializers.xavier_uniform(),\n",
    "               rngs=nnx.Rngs(0)), nnx.relu,\n",
    "    nnx.Linear(8, 1, kernel_init=nnx.initializers.constant(42.0),\n",
    "               rngs=nnx.Rngs(1)))\n",
    "\n",
    "net.layers[0].kernel[:, 0], net.layers[2].kernel"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(Array([0.48410526, 0.63880914, 0.3169195 , 0.4499392 ], dtype=float32),\n Param( # 8 (32 B)\n   value=Array([[42.],\n          [42.],\n          [42.],\n          [42.],\n          [42.],\n          [42.],\n          [42.],\n          [42.]], dtype=float32)\n ))"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-9 fw=jax -->\n",
    "\n",
    "The second route walks the existing object graph. `nnx.iter_modules` yields\n",
    "each child with its path, and one fresh key per matching layer keeps the draws\n",
    "independent. The function below re-draws every linear kernel of the\n",
    "constant-42 network; biases pass through untouched:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-applying-initializers-3"
   },
   "source": [
    "def reinit(model, key, init_fn):\n",
    "    layers = [m for _, m in nnx.iter_modules(model)\n",
    "              if isinstance(m, nnx.Linear)]\n",
    "    for layer, subkey in zip(layers, jax.random.split(key, len(layers))):\n",
    "        layer.kernel[...] = init_fn(\n",
    "            subkey, layer.kernel.shape, layer.kernel.dtype)\n",
    "\n",
    "reinit(net, jax.random.key(1), nnx.initializers.xavier_uniform())\n",
    "net.layers[2].kernel[:3]"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[-0.15734486],\n       [ 0.0798452 ],\n       [-0.4529283 ]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-10 fw=all -->\n",
    "\n",
    "## Modern Schemes: Truncation, Depth, and Zeros\n",
    "\n",
    "Xavier and He set the variance of a single layer. The schemes below, standard\n",
    "in transformer codebases, adjust what happens in the distribution's tails,\n",
    "across depth, and at a block's start.\n",
    "\n",
    "### Truncated Normals\n",
    "\n",
    "A Gaussian gets the variance right, but its tails are unbounded. That is\n",
    "harmless for one draw and a near-certainty at scale: among the $10^8$ weights\n",
    "of a BERT-sized model, dozens land beyond five standard deviations. Large\n",
    "draws also consume disproportionate dynamic range once a model is cast to low\n",
    "precision (that section). BERT and implementations in the ViT\n",
    "lineage use truncated-normal initialization\n",
    "[@Devlin.Chang.Lee.ea.2018; @Dosovitskiy.Beyer.Kolesnikov.ea.2021]: the\n",
    "tails are cut off at a fixed multiple of the nominal standard deviation. Raw\n",
    "truncated-normal initializers do not necessarily restore the variance removed\n",
    "with the tails; fan-aware variance-scaling initializers often do."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-11 fw=jax -->\n",
    "\n",
    "`nnx.initializers.truncated_normal` states its cutoffs in units of the\n",
    "standard deviation (`lower=-2.0`, `upper=2.0` by default), so the clip at two\n",
    "standard deviations is what you get with no extra arguments. Truncation is\n",
    "the house preference throughout Flax: the `variance_scaling` factory behind\n",
    "`lecun_normal`, `xavier_normal`, and `he_normal` draws its normal variants\n",
    "truncated as well."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-truncated-normals"
   },
   "source": [
    "key = jax.random.key(0)\n",
    "w = nnx.initializers.normal(stddev=0.02)(key, (1000, 1000))\n",
    "print(f'normal:    std {w.std():.4f}, max weight {jnp.abs(w).max():.4f}')\n",
    "w = nnx.initializers.truncated_normal(stddev=0.02)(key, (1000, 1000))\n",
    "print(f'truncated: std {w.std():.4f}, max weight {jnp.abs(w).max():.4f}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "normal:    std 0.0200, max weight 0.0929\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "truncated: std 0.0176, max weight 0.0400\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-12 fw=all -->\n",
    "\n",
    "A million plain-normal draws produce entries near $5\\sigma = 0.1$; the\n",
    "truncated version guarantees $|w| \\leq 0.04$. Its printed standard deviation\n",
    "dips slightly below the nominal 0.02 because truncation removes tail mass;\n",
    "practice ignores the difference.\n",
    "\n",
    "### Scaling Down Residual Branches\n",
    "\n",
    "A residual network computes $\\mathbf{x}_{k+1} = \\mathbf{x}_k +\n",
    "f_k(\\mathbf{x}_k)$, so its output is the input plus the sum of $N$ block\n",
    "contributions. If every block is initialized identically, each contribution\n",
    "has variance proportional to that of the already-inflated stream it reads, and\n",
    "the stream's variance compounds geometrically with depth. This is the\n",
    "depth-wise cousin of the layer-wise explosion in\n",
    "that section. GPT-2's fix is to shrink only the *last*\n",
    "linear layer of each residual branch, the output projection that writes into\n",
    "the stream, by a factor of $1/\\sqrt{N}$: with $N$ roughly independent\n",
    "contributions each scaled down that way, the sum's variance stays $O(1)$\n",
    "regardless of depth [@Radford.Wu.Child.ea.2019]. For an $L$-layer\n",
    "transformer the published recipe is a normal with std $0.02/\\sqrt{2L}$ on the\n",
    "residual projections — the number of contributions is $N = 2L$, because each\n",
    "layer writes to the stream twice, once from attention and once from the MLP.\n",
    "Unlike Xavier or He, the base 0.02 is not derived from fan-in; it is an\n",
    "empirical constant that worked at GPT-2's widths and stuck.\n",
    "\n",
    "the figure draws the two regimes side by side.\n",
    "\n",
    "![A residual stream with additive block contributions, unscaled (left) versus scaled by 1/sqrt(N) (right). Under the independent-contribution approximation, unscaled variance grows with N and scaling keeps the sum O(1). In an unnormalized stack, each branch reads an already inflated stream, so growth can be faster, as the experiment below shows.](../img/bg-residual-stream.svg)\n",
    "\n",
    "### Starting a Block at Zero\n",
    "\n",
    "Zero-initializing *every* weight is fatal: all units in a layer compute the\n",
    "same output, receive the same gradient, and stay identical forever\n",
    "(that section). Zero-initializing just the *last* layer\n",
    "of a residual block is a different and useful move. The branch then\n",
    "contributes exactly nothing, each block is the identity map, and the network\n",
    "starts as a shallow function whose depth switches on gradually during\n",
    "training. Symmetry is not a problem because the branch's earlier layers keep\n",
    "their random weights. On the first backward pass, the zeroed projection\n",
    "receives a nonzero gradient because its input is nonzero, while the earlier\n",
    "layers receive zero gradient through that projection. After the projection\n",
    "moves off zero, gradient reaches the whole branch. Keep this scheme distinct\n",
    "from the previous one: GPT-2 makes every residual projection *small but\n",
    "nonzero*, whereas zero-init makes one layer *exactly zero* so the block starts\n",
    "as an exact identity.\n",
    "\n",
    "### Watching the Variance Compound\n",
    "\n",
    "Claims about variance at depth are cheap to test. We reuse a compact residual\n",
    "block (it repeats the one from that section) and stack\n",
    "it $N$ deep:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-1"
   },
   "source": [
    "he = nnx.initializers.variance_scaling(2.0, 'fan_in', 'truncated_normal')\n",
    "\n",
    "class ResidualBlock(nnx.Module):\n",
    "    def __init__(self, d, out_init=he, rngs=None):\n",
    "        rngs = nnx.Rngs(0) if rngs is None else rngs\n",
    "        self.fc1 = nnx.Linear(d, 4 * d, kernel_init=he, rngs=rngs)\n",
    "        self.fc2 = nnx.Linear(4 * d, d, kernel_init=out_init, rngs=rngs)\n",
    "\n",
    "    def __call__(self, X):\n",
    "        return X + self.fc2(nnx.relu(self.fc1(X)))"
   ],
   "execution_count": 7,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-13 fw=jax -->\n",
    "\n",
    "Every dense layer gets He initialization, appropriate for the ReLU inside\n",
    "the branch; `variance_scaling(2.0, 'fan_in', 'truncated_normal')` is exactly\n",
    "what `nnx.initializers.he_normal()` expands to. The block takes its output\n",
    "projection's initializer as a constructor argument, so each treatment is a\n",
    "different initializer.\n",
    "Leaving it alone means He again, scaling by $1/\\sqrt{N}$ is a closure over\n",
    "the depth that wraps the same three-argument signature, and zeroing is\n",
    "`nnx.initializers.zeros`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-2"
   },
   "source": [
    "def scaled_he(num_blocks):\n",
    "    def init(key, shape, dtype=jnp.float32):\n",
    "        return he(key, shape, dtype) * num_blocks ** -0.5\n",
    "    return init\n",
    "\n",
    "def output_std(num_blocks, out_init):\n",
    "    rngs = nnx.Rngs(0)\n",
    "    net = nnx.Sequential(*[\n",
    "        ResidualBlock(64, out_init=out_init(num_blocks), rngs=rngs)\n",
    "        for _ in range(num_blocks)])\n",
    "    X = jax.random.normal(jax.random.key(1), (256, 64))\n",
    "    return net(X).std().item()"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-14 fw=all -->\n",
    "\n",
    "One forward pass on unit-variance inputs per depth and treatment:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-watching-the-variance-compound-3"
   },
   "source": [
    "tweaks = {'default': lambda n: he,\n",
    "          'scaled': scaled_he,\n",
    "          'zero': lambda n: nnx.initializers.zeros}\n",
    "print(f'{\"N\":>3}' + ''.join(f'{name:>10}' for name in tweaks))\n",
    "for n in (2, 8, 32):\n",
    "    stds = (output_std(n, tweak) for tweak in tweaks.values())\n",
    "    print(f'{n:>3}' + ''.join(f'{s:>10.3g}' for s in stds))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  N   default    scaled      zero\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  2      2.89      1.93     0.993\n  8      75.3      2.27     0.993\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " 32   3.6e+07      2.59     0.993\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-15 fw=all -->\n",
    "\n",
    "The default column multiplies by a roughly constant factor per block in this\n",
    "unnormalized stack, reaching tens of millions by $N=32$. Such initial\n",
    "activations make optimization impractical. The scaled column stays near a\n",
    "small constant over the depths tested: scaling cuts each branch's variance by\n",
    "a factor of $N$, following the independent-contribution approximation above.\n",
    "The zero column reproduces the input's standard deviation, since every block\n",
    "is the identity. This forward-pass test does not prove how an optimizer will\n",
    "behave, but it catches an unusable initialization before training begins.\n",
    "\n",
    "## Custom Initializers\n",
    "\n",
    "Occasionally the menu has nothing you need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-16 fw=jax -->\n",
    "\n",
    "An initializer is just a function `(key, shape, dtype) -> array`, the\n",
    "signature everything in `nnx.initializers` shares, so writing one is no harder\n",
    "than using one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-17 fw=all -->\n",
    "\n",
    "Suppose, to make the point vividly, we want weights distributed as\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    w \\sim \\begin{cases}\n",
    "        U(5, 10) & \\textrm{ with probability } \\frac{1}{4} \\\\\n",
    "            0    & \\textrm{ with probability } \\frac{1}{2} \\\\\n",
    "        U(-10, -5) & \\textrm{ with probability } \\frac{1}{4}\n",
    "    \\end{cases}\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-18 fw=jax -->\n",
    "\n",
    "Split the key, draw a magnitude from $U(5, 10)$, and draw a factor from\n",
    "$\\{-1, 0, 0, 1\\}$, which zeroes the entry with probability one half and keeps\n",
    "either sign with probability one quarter. Handing the function to\n",
    "`kernel_init` makes it official; `init` gives every layer an independent key:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-custom-initializers-1"
   },
   "source": [
    "def my_init(key, shape, dtype=jnp.float32):\n",
    "    mag_key, sign_key = jax.random.split(key)\n",
    "    mag = jax.random.uniform(mag_key, shape, dtype, minval=5, maxval=10)\n",
    "    sign = jax.random.choice(sign_key, jnp.array([-1., 0., 0., 1.], dtype),\n",
    "                             shape)\n",
    "    return mag * sign\n",
    "\n",
    "net = nnx.Sequential(\n",
    "    nnx.Linear(4, 8, kernel_init=my_init, rngs=nnx.Rngs(0)), nnx.relu,\n",
    "    nnx.Linear(8, 1, rngs=nnx.Rngs(1)))\n",
    "net.layers[0].kernel[:, :2]"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([[-6.4136   ,  8.940486 ],\n       [-6.2257624,  0.       ],\n       [ 0.       , -9.485439 ],\n       [ 0.       ,  9.542628 ]], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-19 fw=jax -->\n",
    "\n",
    "NNX parameters are mutable variables. Assignment updates the value owned by\n",
    "the layer while preserving the `nnx.Param` object that optimizers and\n",
    "checkpoints track. One-off surgery therefore looks like ordinary indexed\n",
    "assignment:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "init-custom-initializers-2"
   },
   "source": [
    "net.layers[0].kernel[...] += 1\n",
    "net.layers[0].kernel[0, 0] = 42\n",
    "net.layers[0].kernel[0]"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([42.       ,  9.940486 ,  1.       ,  1.       ,  1.       ,\n        1.       , -4.9630113,  1.       ], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-20 fw=jax -->\n",
    "\n",
    "There is no ordering caveat: NNX parameters are created in the constructor,\n",
    "so both graph walks and direct assignment can run immediately."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-21 fw=all -->\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-22 fw=jax -->\n",
    "\n",
    "NNX initializes parameters in layer constructors with fan-aware defaults;\n",
    "override them when depth or a paper's recipe demands it. An initializer is a function\n",
    "`(key, shape, dtype) -> array`, handed to a layer as `kernel_init` at\n",
    "construction or run over selected modules in an existing object graph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-23 fw=all -->\n",
    "\n",
    "On top of Xavier and He,\n",
    "transformer-era code truncates normal tails to bound outliers, shrinks each\n",
    "residual output projection by $1/\\sqrt{N}$ so the stream's variance stays\n",
    "$O(1)$ at any depth, and zero-initializes a block's last layer to start it as\n",
    "the identity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-24 fw=jax -->\n",
    "\n",
    "Anything the menu lacks is a few lines of `jax.random` code behind the same\n",
    "three-argument signature."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=init-md-25 fw=all -->\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Instrument the residual stack: record the standard deviation of the\n",
    "   activation after every block (run the stack one block at a time, or\n",
    "   capture per-block activations with the tools of that section)\n",
    "   for the default and scaled treatments at $N=32$, and plot it against\n",
    "   depth. Which curve matches the geometric-growth prediction?\n",
    "1. Zero-initialize *all* layers of every block instead of just the output\n",
    "   projection. The forward pass still returns the input, but what can the\n",
    "   network learn? Work out which parameters receive a nonzero gradient, and\n",
    "   relate the answer to the symmetry-breaking argument of\n",
    "   that section.\n",
    "1. Write an initializer that fills each parameter from a dictionary keyed by\n",
    "   parameter name (walk `net.named_parameters()` in PyTorch, `net.weights`\n",
    "   in TensorFlow, `nnx.iter_modules(net)` in JAX, `net.collect_params()`\n",
    "   in MXNet). You have re-invented part of checkpoint loading,\n",
    "   which that section covers.\n",
    "1. For a normal distribution truncated at $\\pm 2\\sigma$: what fraction of\n",
    "   draws does the clip discard, and by how much does it shrink the standard\n",
    "   deviation? Verify both numbers against the printed output of the truncation\n",
    "   demo above."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
