{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-1 fw=all -->\n",
    "\n",
    "# Utility Functions and Classes\n",
    "\n",
    "\n",
    "This section contains the implementations of utility functions and classes used in this book."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-1"
   },
   "source": [
    "import inspect\n",
    "from IPython import display\n",
    "import collections\n",
    "from d2l import jax as d2l\n",
    "import jax"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-2 fw=all -->\n",
    "\n",
    "Hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-2"
   },
   "source": [
    "@d2l.add_to_class(d2l.HyperParameters)\n",
    "def save_hyperparameters(self, ignore=None):\n",
    "    \"\"\"Save function arguments into class attributes.\"\"\"\n",
    "    ignore = [] if ignore is None else ignore\n",
    "    frame = inspect.currentframe().f_back\n",
    "    _, _, _, local_vars = inspect.getargvalues(frame)\n",
    "    self.hparams = {k:v for k, v in local_vars.items()\n",
    "                    if k not in set(ignore+['self']) and not k.startswith('_')}\n",
    "    for k, v in self.hparams.items():\n",
    "        setattr(self, k, v)"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-3 fw=all -->\n",
    "\n",
    "Progress bar."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-3"
   },
   "source": [
    "\n",
    "import io, os, queue, threading, time\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def draw(self, x, y, label, every_n=1):\n",
    "    \"\"\"Schedule the point (x, y) to be plotted under `label`.\n",
    "\n",
    "    This call is *asynchronous*: it appends the point to an internal queue\n",
    "    and returns immediately, so the (possibly compiled) training loop never\n",
    "    blocks on a device-to-host transfer or on matplotlib. `y` may be a number\n",
    "    or a zero-argument callable returning one; if it is a callable, the\n",
    "    conversion runs on the background drawing thread, not on the caller's.\n",
    "    `every_n` averages each group of `n` raw points into one plotted point.\"\"\"\n",
    "    self._start_drawer()\n",
    "    try:\n",
    "        self._queue.put_nowait((x, y, label, every_n))\n",
    "    except queue.Full:\n",
    "        pass  # drop the point rather than stall the training loop\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _start_drawer(self):\n",
    "    if getattr(self, '_drawer', None) is not None:\n",
    "        return\n",
    "    self._queue = queue.Queue(maxsize=1000)\n",
    "    self._lock = threading.Lock()\n",
    "    self._handle = None\n",
    "    self._last = 0.0\n",
    "    self._running = True\n",
    "    # Emit live frames only outside the headless book build, so an executed\n",
    "    # notebook records exactly one (final) figure in the committed store.\n",
    "    self._live = self.display and os.environ.get('D2L_NB_CAPTURE') != '1'\n",
    "    self._drawer = threading.Thread(target=self._drain, daemon=True)\n",
    "    self._drawer.start()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _drain(self):\n",
    "    while self._running or not self._queue.empty():\n",
    "        try:\n",
    "            batch = [self._queue.get(timeout=0.1)]\n",
    "        except queue.Empty:\n",
    "            continue\n",
    "        while True:  # coalesce everything currently queued\n",
    "            try: batch.append(self._queue.get_nowait())\n",
    "            except queue.Empty: break\n",
    "        for x, y, label, every_n in batch:\n",
    "            self._accumulate(x, y() if callable(y) else y, label, every_n)\n",
    "        if self._live and time.time() - self._last > 0.1:\n",
    "            self._render(thread=True)\n",
    "            self._last = time.time()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _accumulate(self, x, y, label, every_n):\n",
    "    Point = collections.namedtuple('Point', ['x', 'y'])\n",
    "    with self._lock:\n",
    "        if not hasattr(self, 'raw_points'):\n",
    "            self.raw_points = collections.OrderedDict()\n",
    "            self.data = collections.OrderedDict()\n",
    "        if label not in self.raw_points:\n",
    "            self.raw_points[label] = []\n",
    "            self.data[label] = []\n",
    "        points = self.raw_points[label]\n",
    "        points.append(Point(x, float(y)))\n",
    "        if len(points) != every_n:\n",
    "            return\n",
    "        mean = lambda v: sum(v) / len(v)\n",
    "        self.data[label].append(Point(mean([p.x for p in points]),\n",
    "                                      mean([p.y for p in points])))\n",
    "        points.clear()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _plot_lines(self, axes):\n",
    "    with self._lock:\n",
    "        series = [(k, list(v)) for k, v in getattr(self, 'data', {}).items()]\n",
    "    for (k, v), ls, color in zip(series, self.ls, self.colors):\n",
    "        axes.plot([p.x for p in v], [p.y for p in v],\n",
    "                  linestyle=ls, color=color, label=k)\n",
    "    if self.xlim: axes.set_xlim(self.xlim)\n",
    "    if self.ylim: axes.set_ylim(self.ylim)\n",
    "    if not self.xlabel: self.xlabel = self.x\n",
    "    axes.set_xlabel(self.xlabel)\n",
    "    axes.set_ylabel(self.ylabel)\n",
    "    axes.set_xscale(self.xscale)\n",
    "    # A smoothed series exists before it necessarily contains its first point.\n",
    "    # Matplotlib rejects an empty/nonpositive log axis, so keep the temporary\n",
    "    # live frame linear and apply the requested log scale as soon as a positive\n",
    "    # value is available. The final training figure therefore still uses log.\n",
    "    has_positive_y = any(p.y > 0 for _, values in series for p in values)\n",
    "    if self.yscale != 'log' or has_positive_y:\n",
    "        axes.set_yscale(self.yscale)\n",
    "    if series: axes.legend()\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def _render(self, thread=False):\n",
    "    if not self.display:\n",
    "        return\n",
    "    if thread:\n",
    "        # Off the main thread: matplotlib's global pyplot state is not\n",
    "        # thread-safe, so render into a private Agg figure and push a PNG.\n",
    "        from matplotlib.figure import Figure\n",
    "        from matplotlib.backends.backend_agg import FigureCanvasAgg\n",
    "        fig = Figure(figsize=self.figsize)\n",
    "        FigureCanvasAgg(fig)\n",
    "        self._plot_lines(fig.add_subplot(111))\n",
    "        buf = io.BytesIO()\n",
    "        fig.savefig(buf, format='png', bbox_inches='tight')\n",
    "        frame = display.Image(data=buf.getvalue())\n",
    "        if self._handle is None:\n",
    "            self._handle = display.display(frame, display_id=True)\n",
    "        else:\n",
    "            self._handle.update(frame)\n",
    "    else:\n",
    "        # Main thread (final frame): render through pyplot so the output is\n",
    "        # captured exactly like the rest of the book's inline figures.\n",
    "        d2l.use_svg_display()\n",
    "        if self.fig is None:\n",
    "            self.fig = d2l.plt.figure(figsize=self.figsize)\n",
    "        self.fig.clf()\n",
    "        self._plot_lines(self.fig.add_subplot(111))\n",
    "        if self._handle is not None:\n",
    "            self._handle.update(self.fig)         # interactive: reuse live slot\n",
    "        else:\n",
    "            # Display, then a trailing clear_output(wait=True): the inline\n",
    "            # backend's automatic end-of-cell render then *replaces* (rather\n",
    "            # than duplicates) this figure, so exactly one image is captured.\n",
    "            display.display(self.fig)\n",
    "            display.clear_output(wait=True)\n",
    "\n",
    "@d2l.add_to_class(d2l.ProgressBoard)\n",
    "def flush(self):\n",
    "    \"\"\"Wait for every scheduled point to be drawn, then render the final\n",
    "    figure on the calling thread. `Trainer.fit` calls this for you; call it\n",
    "    yourself after a bare `draw` loop.\"\"\"\n",
    "    if getattr(self, '_drawer', None) is not None:\n",
    "        self._running = False\n",
    "        self._drawer.join()\n",
    "        self._drawer = None\n",
    "    self._render(thread=False)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-4 fw=all -->\n",
    "\n",
    "Add FrozenLake environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-5 fw=all -->\n",
    "\n",
    "Create environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-6 fw=all -->\n",
    "\n",
    "Show value function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-7 fw=all -->\n",
    "\n",
    "Show Q function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-8 fw=all -->\n",
    "\n",
    "Trainer\n",
    "\n",
    "Legacy helper functions retained for backward compatibility:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-8"
   },
   "source": [
    "\n",
    "def linreg(X, w, b):\n",
    "    \"\"\"The linear regression model.\"\"\"\n",
    "    return d2l.matmul(X, w) + b\n",
    "\n",
    "def squared_loss(y_hat, y):\n",
    "    \"\"\"Squared loss.\"\"\"\n",
    "    return (y_hat - d2l.reshape(y, y_hat.shape)) ** 2 / 2\n",
    "\n",
    "def load_array(data_arrays, batch_size, is_train=True):\n",
    "    \"\"\"Construct a JAX-compatible data iterator.\n",
    "\n",
    "    For training, the last partial batch is dropped so every minibatch\n",
    "    has the same shape — this lets a `@jax.jit`'d step function compile\n",
    "    once per epoch instead of recompiling for the smaller last batch.\n",
    "    Validation iterators yield all batches (the last may be smaller),\n",
    "    matching how PyTorch / TF / MXNet behave.\n",
    "    \"\"\"\n",
    "    # Keep dataset storage on the host. Only minibatches need to be transferred\n",
    "    # to an accelerator, and host-side indexing avoids launching a device gather\n",
    "    # for every array in every minibatch.\n",
    "    data_arrays = tuple(np.asarray(a) for a in data_arrays)\n",
    "    n = data_arrays[0].shape[0]\n",
    "    indices = np.arange(n)\n",
    "    last = n - (n % batch_size) if is_train else n\n",
    "    def data_iter():\n",
    "        if is_train:\n",
    "            np.random.shuffle(indices)\n",
    "            # Shuffle each full field once, then transfer contiguous slices.\n",
    "            # This avoids a separate fancy-index gather for every field and\n",
    "            # every minibatch.\n",
    "            epoch_arrays = tuple(a[indices] for a in data_arrays)\n",
    "        else:\n",
    "            epoch_arrays = data_arrays\n",
    "        for i in range(0, last, batch_size):\n",
    "            yield tuple(jnp.array(a[i: min(i + batch_size, n)])\n",
    "                        for a in epoch_arrays)\n",
    "    class DataIter:\n",
    "        def __iter__(self):\n",
    "            return data_iter()\n",
    "        def __len__(self):\n",
    "            return last // batch_size if is_train else (n + batch_size - 1) // batch_size\n",
    "    return DataIter()\n",
    "\n",
    "class ArrayDataLoader:\n",
    "    \"\"\"A simple data loader for JAX that batches arrays or dataset objects.\"\"\"\n",
    "    def __init__(self, *args, batch_size=32, shuffle=False, drop_last=False,\n",
    "                 **kwargs):\n",
    "        if len(args) == 1 and hasattr(args[0], '__len__') and hasattr(args[0], '__getitem__'):\n",
    "            dataset = args[0]\n",
    "            items = [dataset[i] for i in range(len(dataset))]\n",
    "            if isinstance(items[0], (tuple, list)):\n",
    "                self.arrays = tuple(np.array([item[j] for item in items])\n",
    "                                    for j in range(len(items[0])))\n",
    "            else:\n",
    "                self.arrays = (np.array(items),)\n",
    "        else:\n",
    "            self.arrays = tuple(np.asarray(a) for a in args)\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.drop_last = drop_last\n",
    "\n",
    "    def __iter__(self):\n",
    "        n = self.arrays[0].shape[0]\n",
    "        indices = np.arange(n)\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(indices)\n",
    "        for i in range(0, n, self.batch_size):\n",
    "            end = i + self.batch_size\n",
    "            if self.drop_last and end > n:\n",
    "                break\n",
    "            batch_idx = indices[i:min(end, n)]\n",
    "            yield tuple(jnp.array(a[batch_idx]) for a in self.arrays)\n",
    "\n",
    "    def __len__(self):\n",
    "        n = self.arrays[0].shape[0]\n",
    "        if self.drop_last:\n",
    "            return n // self.batch_size\n",
    "        return (n + self.batch_size - 1) // self.batch_size\n",
    "\n",
    "class Animator:\n",
    "    \"\"\"For plotting data in animation.\"\"\"\n",
    "    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,\n",
    "                 ylim=None, xscale='linear', yscale='linear',\n",
    "                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,\n",
    "                 figsize=(3.5, 2.5)):\n",
    "        if legend is None:\n",
    "            legend = []\n",
    "        d2l.use_svg_display()\n",
    "        self.fig, self.axes = d2l.plt.subplots(nrows, ncols, figsize=figsize)\n",
    "        if nrows * ncols == 1:\n",
    "            self.axes = [self.axes, ]\n",
    "        self.config_axes = lambda: d2l.set_axes(\n",
    "            self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)\n",
    "        self.X, self.Y, self.fmts = None, None, fmts\n",
    "\n",
    "    def add(self, x, y):\n",
    "        if not hasattr(y, \"__len__\"):\n",
    "            y = [y]\n",
    "        n = len(y)\n",
    "        if not hasattr(x, \"__len__\"):\n",
    "            x = [x] * n\n",
    "        if not self.X:\n",
    "            self.X = [[] for _ in range(n)]\n",
    "        if not self.Y:\n",
    "            self.Y = [[] for _ in range(n)]\n",
    "        for i, (a, b) in enumerate(zip(x, y)):\n",
    "            if a is not None and b is not None:\n",
    "                self.X[i].append(a)\n",
    "                self.Y[i].append(b)\n",
    "        self.axes[0].cla()\n",
    "        for x, y, fmt in zip(self.X, self.Y, self.fmts):\n",
    "            self.axes[0].plot(x, y, fmt)\n",
    "        self.config_axes()\n",
    "        display.display(self.fig)\n",
    "        display.clear_output(wait=True)\n",
    "\n",
    "class Accumulator:\n",
    "    \"\"\"For accumulating sums over `n` variables.\"\"\"\n",
    "    def __init__(self, n):\n",
    "        self.data = [0.0] * n\n",
    "\n",
    "    def add(self, *args):\n",
    "        self.data = [a + float(b) for a, b in zip(self.data, args)]\n",
    "\n",
    "    def reset(self):\n",
    "        self.data = [0.0] * len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]\n",
    "\n",
    "def tokenize(lines, token='word'):\n",
    "    \"\"\"Split text lines into word or character tokens.\"\"\"\n",
    "    assert token in ('word', 'char'), 'Unknown token type: ' + token\n",
    "    return [line.split() if token == 'word' else list(line) for line in lines]\n",
    "\n",
    "def truncate_pad(line, num_steps, padding_token):\n",
    "    \"\"\"Truncate or pad sequences.\"\"\"\n",
    "    if len(line) > num_steps:\n",
    "        return line[:num_steps]\n",
    "    return line + [padding_token] * (num_steps - len(line))\n",
    "\n",
    "def download_extract(name, folder=None):\n",
    "    \"\"\"Download and extract a zip/tar file.\"\"\"\n",
    "    fname = download(name)\n",
    "    base_dir = os.path.dirname(fname)\n",
    "    data_dir, ext = os.path.splitext(fname)\n",
    "    target = os.path.join(base_dir, folder) if folder else data_dir\n",
    "    marker = fname + '.extracted'\n",
    "    if os.path.exists(marker):\n",
    "        return target\n",
    "    if ext == '.zip':\n",
    "        fp = zipfile.ZipFile(fname, 'r')\n",
    "    elif ext in ('.tar', '.gz'):\n",
    "        fp = tarfile.open(fname, 'r')\n",
    "    else:\n",
    "        assert False, 'Only zip/tar files can be extracted.'\n",
    "    fp.extractall(base_dir)\n",
    "    open(marker, 'w').close()\n",
    "    return target\n",
    "\n",
    "def evaluate_loss(net, data_iter, loss):\n",
    "    \"\"\"Evaluate the loss of a model on the given dataset.\"\"\"\n",
    "    metric = d2l.Accumulator(2)\n",
    "    for X, y in data_iter:\n",
    "        out = net(X)\n",
    "        y = d2l.reshape(y, out.shape)\n",
    "        l = loss(out, y)\n",
    "        metric.add(d2l.reduce_sum(l), d2l.size(l))\n",
    "    return metric[0] / metric[1]"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-10"
   },
   "source": [
    "def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):\n",
    "    \"\"\"Plot a list of images.\"\"\"\n",
    "    figsize = (num_cols * scale, num_rows * scale)\n",
    "    _, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize)\n",
    "    axes = axes.flatten()\n",
    "    for i, (ax, img) in enumerate(zip(axes, imgs)):\n",
    "        try:\n",
    "            img = d2l.numpy(img)\n",
    "        except:\n",
    "            pass\n",
    "        ax.imshow(img)\n",
    "        ax.axes.get_xaxis().set_visible(False)\n",
    "        ax.axes.get_yaxis().set_visible(False)\n",
    "        if titles:\n",
    "            ax.set_title(titles[i])\n",
    "    return axes"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "utils-utility-functions-and-classes-12"
   },
   "source": [
    "\n",
    "import os\n",
    "import requests\n",
    "import zipfile\n",
    "import tarfile\n",
    "import hashlib\n",
    "\n",
    "def download(url, folder='../data', sha1_hash=None):\n",
    "    \"\"\"Download a file to folder and return the local filepath.\"\"\"\n",
    "    if not url.startswith('http'):\n",
    "        # For back compatibility\n",
    "        url, sha1_hash = DATA_HUB[url]\n",
    "    os.makedirs(folder, exist_ok=True)\n",
    "    fname = os.path.join(folder, url.split('/')[-1])\n",
    "    # Check if hit cache\n",
    "    if os.path.exists(fname) and sha1_hash:\n",
    "        sha1 = hashlib.sha1()\n",
    "        with open(fname, 'rb') as f:\n",
    "            while True:\n",
    "                data = f.read(1048576)\n",
    "                if not data:\n",
    "                    break\n",
    "                sha1.update(data)\n",
    "        if sha1.hexdigest() == sha1_hash:\n",
    "            return fname\n",
    "    # Download\n",
    "    import tempfile\n",
    "    print(f'Downloading {fname} from {url}...')\n",
    "    r = requests.get(url, stream=True, verify=True)\n",
    "    fd, tmp = tempfile.mkstemp(dir=folder)\n",
    "    try:\n",
    "        with os.fdopen(fd, 'wb') as f:\n",
    "            f.write(r.content)\n",
    "        os.replace(tmp, fname)\n",
    "    except BaseException:\n",
    "        os.unlink(tmp)\n",
    "        raise\n",
    "    return fname\n",
    "\n",
    "def extract(filename, folder=None):\n",
    "    \"\"\"Extract a zip/tar file into folder.\"\"\"\n",
    "    import tempfile, shutil\n",
    "    base_dir = os.path.dirname(filename)\n",
    "    _, ext = os.path.splitext(filename)\n",
    "    assert ext in ('.zip', '.tar', '.gz'), 'Only support zip/tar files.'\n",
    "    opener = zipfile.ZipFile if ext == '.zip' else tarfile.open\n",
    "    if folder is None:\n",
    "        folder = base_dir\n",
    "    tmp = tempfile.mkdtemp(dir=folder)\n",
    "    try:\n",
    "        with opener(filename, 'r') as fp:\n",
    "            fp.extractall(tmp)\n",
    "        for name in os.listdir(tmp):\n",
    "            src = os.path.join(tmp, name)\n",
    "            dst = os.path.join(folder, name)\n",
    "            if os.path.isdir(dst):\n",
    "                shutil.rmtree(dst)\n",
    "            os.rename(src, dst)\n",
    "    finally:\n",
    "        shutil.rmtree(tmp, ignore_errors=True)"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-9 fw=all -->\n",
    "\n",
    "More for the attention chapter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=utils-md-10 fw=all -->\n",
    "\n",
    "## A Note on Framework Coverage\n",
    "\n",
    "The legacy helpers in this section (`evaluate_accuracy`, `train_ch6`,\n",
    "`train_seq2seq`, `predict_seq2seq`, `MaskedSoftmaxCELoss`) are kept for\n",
    "parity with the original D2L implementation, which predates the unified\n",
    "`d2l.Trainer` class introduced in this edition. They are deliberately\n",
    "not provided for JAX, and PyTorch only ships the subset that is genuinely\n",
    "useful outside of the `Trainer` flow. If you are reading the JAX tab,\n",
    "the corresponding chapters use `d2l.Trainer.fit(model, data)` end-to-end;\n",
    "the per-batch logic that these helpers spell out lives inside\n",
    "`Trainer.fit_epoch` and the `@d2l.add_to_class` extensions defined\n",
    "alongside it. The MXNet and TensorFlow tabs additionally retain\n",
    "`evaluate_accuracy` because some of their earlier-chapter snippets call\n",
    "it directly; for PyTorch and JAX,\n",
    "use `Trainer.validation_step` (or its return value) instead."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
