{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-1 fw=all -->\n",
    "\n",
    "# Custom Layers and Functions\n",
    "\n",
    "The library ships over a hundred layers, yet every one of them started life as\n",
    "code somebody wrote because the layer they needed did not exist. Sooner or\n",
    "later you will be in the same position: a new normalization, an unusual\n",
    "residual block, an operation with no gradient. This section shows what to do.\n",
    "A custom layer is a subclass of the module class from\n",
    "that section with a forward method, and if you\n",
    "register its state properly you inherit everything a built-in layer gets:\n",
    "parameter tracking, serialization, and device movement, with no extra code. We\n",
    "build up from a stateless five-liner to RMSNorm, a normalization used by many\n",
    "current language models, then to layers with precomputed non-trainable state,\n",
    "and finally to the case where the forward computation alone is not enough\n",
    "because the gradient itself must be redefined."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-custom-layers-and-functions"
   },
   "source": [
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx\n",
    "from flax import serialization\n",
    "from d2l import jax as d2l"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-2 fw=all -->\n",
    "\n",
    "## Layers without Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-3 fw=jax -->\n",
    "\n",
    "The smallest custom layer has no state at all. `CenteredLayer` subtracts the\n",
    "mean from its input. To build it, we inherit from the base module class and\n",
    "implement `__call__`. This layer owns no variables and needs no constructor."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-layers-without-parameters-1"
   },
   "source": [
    "class CenteredLayer(nnx.Module):\n",
    "    def __call__(self, X):\n",
    "        return X - X.mean()"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-4 fw=all -->\n",
    "\n",
    "Feeding data through confirms that it does what it says."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-5 fw=jax -->\n",
    "\n",
    "NNX modules are ordinary Python objects: calling the object invokes\n",
    "`__call__` directly. Since `CenteredLayer` owns no variables, construction\n",
    "needs neither an input shape nor an RNG stream."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-layers-without-parameters-2"
   },
   "source": [
    "layer = CenteredLayer()\n",
    "layer(jnp.array([1.0, 2, 3, 4, 5]))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([-2., -1.,  0.,  1.,  2.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 3
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-6 fw=all -->\n",
    "\n",
    "Nothing distinguishes this class from a built-in layer. We can place it inside\n",
    "a `Sequential`, and the container neither knows nor cares that one of its\n",
    "children is user code. The output mean should be zero; because we are adding\n",
    "up floating-point numbers, we may see a very small nonzero value instead,\n",
    "which is roundoff, not a bug."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-7 fw=jax -->\n",
    "\n",
    "NNX creates the `Linear` parameters in its constructor, so we specify both\n",
    "feature dimensions there. `CenteredLayer` contributes no parameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-layers-without-parameters-3"
   },
   "source": [
    "net = nnx.Sequential(nnx.Linear(8, 128, rngs=nnx.Rngs(0)), CenteredLayer())\n",
    "Y = net(jax.random.uniform(d2l.get_key(), (4, 8)))\n",
    "Y.mean()"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array(-2.3283064e-09, dtype=float32)"
     },
     "metadata": {},
     "execution_count": 4
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-8 fw=all -->\n",
    "\n",
    "## Layers with Parameters: RMSNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-9 fw=jax -->\n",
    "\n",
    "A layer with something to learn must create its own parameters, and\n",
    "`nnx.Param` is what registers them in the object graph\n",
    "(that section). We could show the mechanics by re-implementing\n",
    "`nnx.Linear`, but that teaches nothing the built-in does not already do. Instead\n",
    "we implement *RMSNorm* [@Zhang.Sennrich.2019], a normalization used by\n",
    "many current language models, in the same handful of lines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-10 fw=all -->\n",
    "\n",
    "Layer normalization standardizes each input vector: subtract the mean, divide\n",
    "by the standard deviation, then apply a learned scale and shift. Zhang and\n",
    "Sennrich observed that the re-centering contributes little and dropped it,\n",
    "along with the shift. What remains is: divide by the root mean square,\n",
    "multiply by a learned gain $\\mathbf{g}$:\n",
    "\n",
    "$$\n",
    "\\textrm{RMSNorm}(\\mathbf{x}) = \\frac{\\mathbf{x}}{\\sqrt{\\frac{1}{d} \\sum_{i=1}^{d} x_i^2 + \\epsilon}} \\odot \\mathbf{g},\n",
    "$$\n",
    "\n",
    "where $d$ is the width of $\\mathbf{x}$ and $\\epsilon$ guards against division\n",
    "by zero. One parameter vector, one reduction, no shift. Why dropping the mean\n",
    "costs so little in practice is a question for the normalization discussion in\n",
    "later chapters; here it is our stateful worked example."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-layers-with-parameters-rmsnorm-1"
   },
   "source": [
    "class RMSNorm(nnx.Module):\n",
    "    def __init__(self, d, eps=1e-6):\n",
    "        self.gain = nnx.Param(jnp.ones(d))\n",
    "        self.eps = eps\n",
    "\n",
    "    def __call__(self, X):\n",
    "        rms = jnp.sqrt((X ** 2).mean(-1, keepdims=True) + self.eps)\n",
    "        return self.gain * X / rms"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-11 fw=all -->\n",
    "\n",
    "Let $m_2$ denote an input row's mean square. With gain one, the output mean\n",
    "square is $m_2/(m_2 + \\epsilon)$. It is therefore close to one when\n",
    "$m_2 \\gg \\epsilon$, as the badly scaled inputs below verify."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-layers-with-parameters-rmsnorm-2"
   },
   "source": [
    "norm = RMSNorm(8)\n",
    "X = 100 * jax.random.normal(d2l.get_key(), (4, 8))\n",
    "(norm(X) ** 2).mean(-1)"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([1.       , 0.9999998, 1.0000001, 1.0000001], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-12 fw=all -->\n",
    "\n",
    "### The Composability Guarantee\n",
    "\n",
    "The reason to write RMSNorm as a module, rather than as a function with a\n",
    "gain tensor floating around beside it, is what registration buys. A correctly\n",
    "written custom layer is indistinguishable from a built-in one along four\n",
    "axes: its parameters are tracked, it composes inside containers, its state\n",
    "serializes, and it moves across devices. We check each once."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-13 fw=jax -->\n",
    "\n",
    "First, assigning `nnx.Param` registered the gain in the NNX object graph.\n",
    "`nnx.state(norm, nnx.Param)` selects it, and an `nnx.Optimizer` configured\n",
    "with `wrt=nnx.Param` will update it."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-the-composability-guarantee-1"
   },
   "source": [
    "nnx.state(norm, nnx.Param)"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "State({\n  'gain': Param( # 8 (32 B)\n    value=Array([1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)\n  )\n})"
     },
     "metadata": {},
     "execution_count": 7
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-14 fw=all -->\n",
    "\n",
    "Second, the layer drops into a `Sequential` next to built-ins."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-the-composability-guarantee-2"
   },
   "source": [
    "net = nnx.Sequential(nnx.Linear(20, 8, rngs=nnx.Rngs(0)), RMSNorm(8),\n",
    "                     nnx.Linear(8, 2, rngs=nnx.Rngs(1)))\n",
    "Y = net(jax.random.normal(d2l.get_key(), (4, 20)))\n",
    "Y.shape"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(4, 2)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-15 fw=jax -->\n",
    "\n",
    "Third, its state serializes with everything else. `flax.serialization` turns\n",
    "the pure dictionary view of NNX state into bytes and back; restoring those\n",
    "values into the state of a fresh model makes the two models agree exactly\n",
    "(saving to disk works the same way; see that section)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-the-composability-guarantee-3"
   },
   "source": [
    "graph, state = nnx.split(net)\n",
    "raw = serialization.to_bytes(nnx.to_pure_dict(state))\n",
    "X = jax.random.normal(d2l.get_key(), (4, 20))\n",
    "clone = nnx.Sequential(nnx.Linear(20, 8, rngs=nnx.Rngs(2)), RMSNorm(8),\n",
    "                       nnx.Linear(8, 2, rngs=nnx.Rngs(3)))\n",
    "clone_graph, clone_state = nnx.split(clone)\n",
    "restored = serialization.from_bytes(nnx.to_pure_dict(clone_state), raw)\n",
    "nnx.replace_by_pure_dict(clone_state, restored)\n",
    "clone = nnx.merge(clone_graph, clone_state)\n",
    "bool(jnp.array_equal(net(X), clone(X)))"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 9
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-16 fw=jax -->\n",
    "\n",
    "Fourth, it moves. NNX state is a pytree of `Variable` objects whose array\n",
    "values can be placed on a chosen device and written back with `nnx.update`.\n",
    "On a machine with a GPU (JAX's default device there) the cell below reports a\n",
    "GPU device twice; on a CPU-only machine it reports `cpu` and runs just the\n",
    "same."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-the-composability-guarantee-4"
   },
   "source": [
    "device = jax.devices()[0]\n",
    "nnx.update(net, jax.device_put(nnx.state(net), device))\n",
    "net.layers[1].gain.device, net(X).device"
   ],
   "execution_count": 10,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(CudaDevice(id=0), CudaDevice(id=0))"
     },
     "metadata": {},
     "execution_count": 10
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-17 fw=all -->\n",
    "\n",
    "None of this took any code beyond the class definition. The guarantee comes\n",
    "from the base class: subclass it, register state through the proper channels,\n",
    "and containers, optimizers, checkpoints, and devices all treat your layer as\n",
    "native.\n",
    "\n",
    "### Checking against the Built-in\n",
    "<!-- d2l:prose id=custom-layers-md-18 fw=pytorch, jax, tensorflow -->\n",
    "\n",
    "RMSNorm proved useful enough that the library now ships its own\n",
    "implementation. That gives us a referee. We copy a nontrivial gain into both\n",
    "implementations, so that agreement cannot be an accident of default values,\n",
    "and compare outputs."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-checking-against-the-built-in"
   },
   "source": [
    "ours = RMSNorm(8)\n",
    "native = nnx.RMSNorm(8, epsilon=1e-6, rngs=nnx.Rngs(0))\n",
    "gain = jnp.linspace(0.5, 1.5, 8)\n",
    "X = jax.random.normal(d2l.get_key(), (16, 8))\n",
    "ours.gain[...] = gain\n",
    "native.scale[...] = gain\n",
    "bool(jnp.allclose(ours(X), native(X)))"
   ],
   "execution_count": 11,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "True"
     },
     "metadata": {},
     "execution_count": 11
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-20 fw=all -->\n",
    "\n",
    "## Precomputed State: Buffers\n",
    "\n",
    "Some layer state is neither a parameter nor a passing activation: a causal\n",
    "mask, a fixed positional table, running statistics. Such values must\n",
    "persist, serialize, and move to the device together with the model — yet no\n",
    "optimizer may ever touch them. Frameworks give this third kind of state its\n",
    "own channel, and custom layers are where you create it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-21 fw=jax -->\n",
    "\n",
    "that section introduced state that persists and travels with\n",
    "the model but receives no gradient. NNX distinguishes kinds of state through\n",
    "`Variable` subclasses. `nnx.Param` marks trainable state; a custom subclass\n",
    "can mark a buffer that an optimizer must skip. A common example is the causal\n",
    "mask that keeps attention scores from looking at future positions. The mask\n",
    "is fixed, so a parameter is wrong. Storing it in `Buffer` keeps it out of the\n",
    "gradient computation while it still serializes and moves with the model.\n",
    "BatchNorm uses the same idea for its running statistics through\n",
    "`nnx.BatchStat`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-22 fw=all -->\n",
    "\n",
    "One caveat before the code: the compact mask below accepts square\n",
    "self-attention score matrices. Cached decoding, where query and key lengths\n",
    "differ, needs an offset mask rather than this $T \\times T$ slice."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-precomputed-state-buffers-1"
   },
   "source": [
    "class Buffer(nnx.Variable):\n",
    "    pass\n",
    "\n",
    "class CausalMask(nnx.Module):\n",
    "    def __init__(self, max_len):\n",
    "        self.mask = Buffer(jnp.triu(\n",
    "            jnp.ones((max_len, max_len), dtype=bool), 1))\n",
    "\n",
    "    def __call__(self, scores):\n",
    "        T = scores.shape[-1]\n",
    "        return jnp.where(self.mask[:T, :T], -jnp.inf, scores)"
   ],
   "execution_count": 12,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-23 fw=jax -->\n",
    "\n",
    "The layer masks the strict upper triangle, has no `nnx.Param` variables for\n",
    "the optimizer, and still carries the mask as `Buffer` state."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-precomputed-state-buffers-2"
   },
   "source": [
    "mask = CausalMask(max_len=8)\n",
    "print(mask(jnp.zeros((3, 3))))\n",
    "print(nnx.state(mask, nnx.Param), nnx.state(mask, Buffer))"
   ],
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[[  0. -inf -inf]\n [  0.   0. -inf]\n [  0.   0.   0.]]\n\u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;255;213;3m})\u001b[0m \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n  \u001b[38;2;156;220;254m'mask'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mBuffer\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 64 (64 B)\u001b[0m\n    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[False,  True,  True,  True,  True,  True,  True,  True],\n           [False, False,  True,  True,  True,  True,  True,  True],\n           [False, False, False,  True,  True,  True,  True,  True],\n           [False, False, False, False,  True,  True,  True,  True],\n           [False, False, False, False, False,  True,  True,  True],\n           [False, False, False, False, False, False,  True,  True],\n           [False, False, False, False, False, False, False,  True],\n           [False, False, False, False, False, False, False, False]],      dtype=bool)\n  \u001b[38;2;255;213;3m)\u001b[0m\n\u001b[38;2;255;213;3m})\u001b[0m\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-24 fw=all -->\n",
    "\n",
    "## Custom Gradients\n",
    "\n",
    "So far we customized the forward computation and let automatic\n",
    "differentiation derive the backward pass. That derivation is literal: it\n",
    "differentiates exactly the operations you ran. Occasionally literalness is\n",
    "the problem. Quantization rounds values to a grid, and rounding is flat\n",
    "almost everywhere, so its true derivative is zero. Any parameter sitting\n",
    "behind a rounding operation stops learning:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-custom-gradients-1"
   },
   "source": [
    "w = jnp.array([0.9, 1.4, 2.6])\n",
    "jax.grad(lambda w: jnp.round(w).sum())(w)"
   ],
   "execution_count": 14,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([0., 0., 0.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 14
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-25 fw=all -->\n",
    "\n",
    "Zero gradient, exactly as calculus demands, and useless for training. The\n",
    "*straight-through estimator* [@Bengio.Leonard.Courville.2013] resolves\n",
    "this with a controlled lie: keep the rounding in the forward pass, but\n",
    "pretend in the backward pass that it was the identity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-26 fw=jax -->\n",
    "\n",
    "No automatic system can derive a lie for us, so we override the chain rule\n",
    "with `jax.custom_vjp`, attaching a hand-written backward rule (a\n",
    "vector-Jacobian product) to an ordinary function, the split\n",
    "the figure draws explicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-27 fw=all -->\n",
    "\n",
    "![The straight-through estimator, forward versus backward. Forward keeps the true staircase round(x), close to but not the same as the identity it approximates; backward substitutes a constant surrogate gradient of 1, the identity's own derivative, for the true gradient, which is zero almost everywhere and would stop all learning.](../img/bg-ste.svg)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-custom-gradients-2"
   },
   "source": [
    "@jax.custom_vjp\n",
    "def round_ste(X):\n",
    "    return jnp.round(X)\n",
    "\n",
    "def round_ste_fwd(X):\n",
    "    return jnp.round(X), None  # (output, residuals); we need no residuals\n",
    "\n",
    "def round_ste_bwd(residuals, grad_output):\n",
    "    # Pretend forward was the identity: pass the gradient straight through\n",
    "    return (grad_output,)\n",
    "\n",
    "round_ste.defvjp(round_ste_fwd, round_ste_bwd)"
   ],
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-28 fw=jax -->\n",
    "\n",
    "Two rules govern the definition. First, the forward rule must return a\n",
    "*pair*: the output plus whatever residuals the backward rule will need. Ours\n",
    "needs nothing and returns `None`; returning the output alone is the most\n",
    "common first-time bug, and JAX reports it as a puzzling structure mismatch\n",
    "rather than pointing at the missing residuals. Second, the backward rule\n",
    "receives those residuals and the loss gradient with respect to the output,\n",
    "and must return a tuple holding one gradient per argument of the original\n",
    "function, hence `(grad_output,)` rather than `grad_output`. An argument that\n",
    "is not differentiable data, say a configuration flag, must be declared\n",
    "through `nondiff_argnums` so it is routed past the gradient machinery."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-29 fw=all -->\n",
    "\n",
    "A small example verifies that the gradient now flows. We multiply the rounded\n",
    "values by known weights, so we know what gradient to expect at `w`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-custom-gradients-3"
   },
   "source": [
    "jax.grad(lambda w: (round_ste(w) * jnp.array([1.0, 2.0, 3.0])).sum())(w)"
   ],
   "execution_count": 16,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Array([1., 2., 3.], dtype=float32)"
     },
     "metadata": {},
     "execution_count": 16
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-30 fw=all -->\n",
    "\n",
    "The downstream gradient arrives at `w` untouched, as if the rounding were the\n",
    "identity, while the forward pass still computed with rounded values. Some\n",
    "estimators instead zero the gradient where a saved input had saturated, while\n",
    "ordinary gradient clipping bounds a large upstream gradient independently of\n",
    "the quantizer. These are different approximations. Neither is a universal\n",
    "default; the surrogate backward rule must match the training objective it is\n",
    "meant to approximate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-31 fw=jax -->\n",
    "\n",
    "JAX offers a shortcut for this particular lie. `jax.lax.stop_gradient` is an\n",
    "identity whose gradient is zero, so `X + stop_gradient(round(X) - X)`\n",
    "computes `round(X)` in the forward pass while the backward pass sees only the\n",
    "leading `X`. Three lines, no `custom_vjp`; the general mechanism earns its\n",
    "keep when the surrogate is not the identity, an input-dependent saturation\n",
    "rule for instance. The two definitions agree in value and in gradient:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "custom-layers-custom-gradients-4"
   },
   "source": [
    "def round_ste2(X):\n",
    "    return X + jax.lax.stop_gradient(jnp.round(X) - X)\n",
    "\n",
    "def loss(f, w):\n",
    "    return (f(w) * jnp.array([1.0, 2.0, 3.0])).sum()\n",
    "\n",
    "(bool(jnp.allclose(round_ste(w), round_ste2(w))),\n",
    " bool(jnp.allclose(jax.grad(loss, 1)(round_ste, w),\n",
    "                   jax.grad(loss, 1)(round_ste2, w))))"
   ],
   "execution_count": 17,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(True, True)"
     },
     "metadata": {},
     "execution_count": 17
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-32 fw=all -->\n",
    "\n",
    "A scope note to close. A custom gradient redefines the gradient of a\n",
    "computation assembled from existing tensor operations; it does not create new\n",
    "ones. Writing new *kernels*, code that runs on the accelerator itself, is a\n",
    "separate craft that we do not cover in this book; that section\n",
    "discusses how to get performance out of the operations you already have.\n",
    "\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-33 fw=jax -->\n",
    "\n",
    "A custom layer is a module subclass: `__call__` defines the computation,\n",
    "`nnx.Param` registers learnable state, and another `nnx.Variable` subclass\n",
    "registers persistent state that the optimizer does not touch. Registration\n",
    "is what buys composability; a correctly written layer\n",
    "gets parameter tracking, container compatibility, serialization, and device\n",
    "movement for free, as we verified on RMSNorm axis by axis. When the chain\n",
    "rule itself must be overridden, as in the straight-through estimator,\n",
    "`jax.custom_vjp` attaches a forward and a backward rule to a function, and\n",
    "`jax.lax.stop_gradient` covers the identity-surrogate case in a single\n",
    "expression. Build custom implementations to understand them; prefer the\n",
    "native ones in production."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-34 fw=all -->\n",
    "\n",
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=custom-layers-md-35 fw=jax -->\n",
    "\n",
    "1. Add an optional learned bias to `RMSNorm` (a shift applied after the\n",
    "   scale, restoring part of what LayerNorm had). Verify that the parameter\n",
    "   state of a model containing it grows by the expected entry, and observe\n",
    "   what `flax.serialization.from_bytes` does when it restores bytes saved\n",
    "   without the bias into the new state structure.\n",
    "1. Implement `Dropout` from scratch as a custom layer that zeroes each entry\n",
    "   with probability $p$ and rescales the survivors during training, but is\n",
    "   the identity during evaluation. Store `self.deterministic = False`, give\n",
    "   the module an `nnx.Rngs` stream, and implement\n",
    "   `set_view(self, *, deterministic)` to update the flag. Verify that\n",
    "   `nnx.view(model, deterministic=True)` reaches a dropout layer nested in an\n",
    "   `nnx.Sequential`.\n",
    "1. Implement a clamp with a custom gradient: a `custom_vjp` function whose\n",
    "   forward is `jnp.clip(X, lo, hi)` and whose backward passes the gradient\n",
    "   only where the input lay strictly inside the clamp range; the forward rule\n",
    "   must save `X` as a residual. Compare its gradients with those of the\n",
    "   native `jnp.clip` for inputs inside, outside, and exactly on the boundary.\n",
    "   Which convention does the native operation use at the boundary?\n",
    "1. Write a layer that returns the leading half of the Fourier coefficients of\n",
    "   its input. It has no parameters, so nothing registers. What do you still\n",
    "   gain by wrapping the computation in a module instead of calling the\n",
    "   transform inline wherever you need it?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
