notebooks-d2l/lstm_jax.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks-d2l/lstm_jax.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9Oi4rCLLlDD"
      },
      "source": [
        "# Long short term memory (LSTM) \n",
        "\n",
        "We show how to implement LSTMs from scratch.\n",
        "Based on sec 9.2 of http://d2l.ai/chapter_recurrent-modern/lstm.html.\n",
        "This uses code from the [basic RNN colab](https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/rnn_torch.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RlKTh9x6ohkD",
        "outputId": "102d19ac-5eaf-40fd-cd9a-e2cfb045bb98"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l\r\u001b[K     |█▉                              | 10 kB 15.2 MB/s eta 0:00:01\r\u001b[K     |███▋                            | 20 kB 20.6 MB/s eta 0:00:01\r\u001b[K     |█████▍                          | 30 kB 24.6 MB/s eta 0:00:01\r\u001b[K     |███████▏                        | 40 kB 16.6 MB/s eta 0:00:01\r\u001b[K     |█████████                       | 51 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |██████████▊                     | 61 kB 15.6 MB/s eta 0:00:01\r\u001b[K     |████████████▌                   | 71 kB 12.7 MB/s eta 0:00:01\r\u001b[K     |██████████████▎                 | 81 kB 11.9 MB/s eta 0:00:01\r\u001b[K     |████████████████                | 92 kB 13.1 MB/s eta 0:00:01\r\u001b[K     |█████████████████▉              | 102 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |███████████████████▋            | 112 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |█████████████████████▍          | 122 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |███████████████████████▏        | 133 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████       | 143 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▊     | 153 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▌   | 163 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▎ | 174 kB 13.5 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 184 kB 13.5 MB/s \n",
            "\u001b[K     |████████████████████████████████| 136 kB 44.7 MB/s \n",
            "\u001b[K     |████████████████████████████████| 72 kB 685 kB/s \n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip install -q flax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "sFsATGYz00TY"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import math\n",
        "from IPython import display \n",
        "\n",
        "import jax\n",
        "import flax.linen as nn\n",
        "import optax\n",
        "\n",
        "import collections\n",
        "import re\n",
        "import random\n",
        "import os\n",
        "import requests\n",
        "import hashlib\n",
        "import time\n",
        "import functools\n",
        "\n",
        "random.seed(0)\n",
        "rng = jax.random.PRNGKey(0)\n",
        "\n",
        "!mkdir figures # for saving plots"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__8VLwdEL2BX"
      },
      "source": [
        "# Data\n",
        "\n",
        " As data, we use the book \"The Time Machine\" by H G Wells,\n",
        "preprocessed using the code in [this colab](https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/text_preproc_torch.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "tLIKRJomvBV5"
      },
      "outputs": [],
      "source": [
        "class SeqDataLoader:\n",
        "    \"\"\"An iterator to load sequence data.\"\"\"\n",
        "    def __init__(self, batch_size, num_steps, use_random_iter, max_tokens):\n",
        "        if use_random_iter:\n",
        "            self.data_iter_fn = seq_data_iter_random\n",
        "        else:\n",
        "            self.data_iter_fn = seq_data_iter_sequential\n",
        "        self.corpus, self.vocab = load_corpus_time_machine(max_tokens)\n",
        "        self.batch_size, self.num_steps = batch_size, num_steps\n",
        "\n",
        "    def __iter__(self):\n",
        "        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)\n",
        "\n",
        "class Vocab:\n",
        "    \"\"\"Vocabulary for text.\"\"\"\n",
        "    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):\n",
        "        if tokens is None:\n",
        "            tokens = []\n",
        "        if reserved_tokens is None:\n",
        "            reserved_tokens = []\n",
        "        # Sort according to frequencies\n",
        "        counter = count_corpus(tokens)\n",
        "        self.token_freqs = sorted(counter.items(), key=lambda x: x[1],\n",
        "                                  reverse=True)\n",
        "        # The index for the unknown token is 0\n",
        "        self.unk, uniq_tokens = 0, ['<unk>'] + reserved_tokens\n",
        "        uniq_tokens += [\n",
        "            token for token, freq in self.token_freqs\n",
        "            if freq >= min_freq and token not in uniq_tokens]\n",
        "        self.idx_to_token, self.token_to_idx = [], dict()\n",
        "        for token in uniq_tokens:\n",
        "            self.idx_to_token.append(token)\n",
        "            self.token_to_idx[token] = len(self.idx_to_token) - 1\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.idx_to_token)\n",
        "\n",
        "    def __getitem__(self, tokens):\n",
        "        if not isinstance(tokens, (list, tuple)):\n",
        "            return self.token_to_idx.get(tokens, self.unk)\n",
        "        return [self.__getitem__(token) for token in tokens]\n",
        "\n",
        "    def to_tokens(self, indices):\n",
        "        if not isinstance(indices, (list, tuple)):\n",
        "            return self.idx_to_token[indices]\n",
        "        return [self.idx_to_token[index] for index in indices]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "-TvZsL4gtHQ2"
      },
      "outputs": [],
      "source": [
        "def tokenize(lines, token='word'):\n",
        "    \"\"\"Split text lines into word or character tokens.\"\"\"\n",
        "    if token == 'word':\n",
        "        return [line.split() for line in lines]\n",
        "    elif token == 'char':\n",
        "        return [list(line) for line in lines]\n",
        "    else:\n",
        "        print('ERROR: unknown token type: ' + token)\n",
        "\n",
        "def count_corpus(tokens):\n",
        "    \"\"\"Count token frequencies.\"\"\"\n",
        "    # Here `tokens` is a 1D list or 2D list\n",
        "    if len(tokens) == 0 or isinstance(tokens[0], list):\n",
        "        # Flatten a list of token lists into a list of tokens\n",
        "        tokens = [token for line in tokens for token in line]\n",
        "    return collections.Counter(tokens)\n",
        "\n",
        "def seq_data_iter_random(corpus, batch_size, num_steps):\n",
        "    \"\"\"Generate a minibatch of subsequences using random sampling.\"\"\"\n",
        "    # Start with a random offset (inclusive of `num_steps - 1`) to partition a\n",
        "    # sequence\n",
        "    corpus = corpus[random.randint(0, num_steps - 1):]\n",
        "    # Subtract 1 since we need to account for labels\n",
        "    num_subseqs = (len(corpus) - 1) // num_steps\n",
        "    # The starting indices for subsequences of length `num_steps`\n",
        "    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))\n",
        "    # In random sampling, the subsequences from two adjacent random\n",
        "    # minibatches during iteration are not necessarily adjacent on the\n",
        "    # original sequence\n",
        "    random.shuffle(initial_indices)\n",
        "\n",
        "    def data(pos):\n",
        "        # Return a sequence of length `num_steps` starting from `pos`\n",
        "        return corpus[pos:pos + num_steps]\n",
        "\n",
        "    num_batches = num_subseqs // batch_size\n",
        "    for i in range(0, batch_size * num_batches, batch_size):\n",
        "        # Here, `initial_indices` contains randomized starting indices for\n",
        "        # subsequences\n",
        "        initial_indices_per_batch = initial_indices[i:i + batch_size]\n",
        "        X = [data(j) for j in initial_indices_per_batch]\n",
        "        Y = [data(j + 1) for j in initial_indices_per_batch]\n",
        "        yield jnp.array(X), jnp.array(Y)\n",
        "\n",
        "def seq_data_iter_sequential(corpus, batch_size, num_steps):\n",
        "    \"\"\"Generate a minibatch of subsequences using sequential partitioning.\"\"\"\n",
        "    # Start with a random offset to partition a sequence\n",
        "    offset = random.randint(0, num_steps)\n",
        "    num_tokens = ((len(corpus) - offset - 1) // batch_size) * batch_size\n",
        "    Xs = jnp.array(corpus[offset:offset + num_tokens])\n",
        "    Ys = jnp.array(corpus[offset + 1:offset + 1 + num_tokens])\n",
        "    Xs, Ys = Xs.reshape(batch_size, -1), Ys.reshape(batch_size, -1)\n",
        "    num_batches = Xs.shape[1] // num_steps\n",
        "    for i in range(0, num_steps * num_batches, num_steps):\n",
        "        X = Xs[:, i:i + num_steps]\n",
        "        Y = Ys[:, i:i + num_steps]\n",
        "        yield X, Y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "yRp-MH0Nv7rN"
      },
      "outputs": [],
      "source": [
        "def download(name, cache_dir=os.path.join('..', 'data')):\n",
        "    \"\"\"Download a file inserted into DATA_HUB, return the local filename.\"\"\"\n",
        "    assert name in DATA_HUB, f\"{name} does not exist in {DATA_HUB}.\"\n",
        "    url, sha1_hash = DATA_HUB[name]\n",
        "    os.makedirs(cache_dir, exist_ok=True)\n",
        "    fname = os.path.join(cache_dir, url.split('/')[-1])\n",
        "    if os.path.exists(fname):\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  # Hit cache\n",
        "    print(f'Downloading {fname} from {url}...')\n",
        "    r = requests.get(url, stream=True, verify=True)\n",
        "    with open(fname, 'wb') as f:\n",
        "        f.write(r.content)\n",
        "    return fname\n",
        "\n",
        "def read_time_machine():\n",
        "    \"\"\"Load the time machine dataset into a list of text lines.\"\"\"\n",
        "    with open(download('time_machine'), 'r') as f:\n",
        "        lines = f.readlines()\n",
        "    return [re.sub('[^A-Za-z]+', ' ', line).strip().lower() for line in lines]\n",
        "\n",
        "def load_corpus_time_machine(max_tokens=-1):\n",
        "    \"\"\"Return token indices and the vocabulary of the time machine dataset.\"\"\"\n",
        "    lines = read_time_machine()\n",
        "    tokens = tokenize(lines, 'char')\n",
        "    vocab = Vocab(tokens)\n",
        "    # Since each text line in the time machine dataset is not necessarily a\n",
        "    # sentence or a paragraph, flatten all the text lines into a single list\n",
        "    corpus = [vocab[token] for line in tokens for token in line]\n",
        "    if max_tokens > 0:\n",
        "        corpus = corpus[:max_tokens]\n",
        "    return corpus, vocab\n",
        "\n",
        "def load_data_time_machine(batch_size, num_steps, use_random_iter=False,\n",
        "                           max_tokens=10000):\n",
        "    \"\"\"Return the iterator and the vocabulary of the time machine dataset.\"\"\"\n",
        "    data_iter = SeqDataLoader(batch_size, num_steps, use_random_iter,\n",
        "                              max_tokens)\n",
        "    return data_iter, data_iter.vocab"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a-IoZP8ILvTB",
        "outputId": "89b1b159-633d-45eb-b00e-a0a981023c7e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading ../data/timemachine.txt from http://d2l-data.s3-accelerate.amazonaws.com/timemachine.txt...\n"
          ]
        }
      ],
      "source": [
        "DATA_HUB = dict()\n",
        "DATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'\n",
        "DATA_HUB['time_machine'] = (DATA_URL + 'timemachine.txt',\n",
        "                                '090b5e7e70c295757f55df93cb0a180b9691891a')\n",
        "\n",
        "batch_size, num_steps = 32, 35\n",
        "train_iter, vocab = load_data_time_machine(batch_size, num_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lNYqmIFwL8Mg"
      },
      "source": [
        "# Creating model from scratch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "tqpiQdOIL9R8"
      },
      "outputs": [],
      "source": [
        "def get_lstm_params(vocab_size, num_hiddens, init_rng):\n",
        "    num_inputs = num_outputs = vocab_size\n",
        "\n",
        "    def normal(shape, rng):\n",
        "        return jax.random.normal(rng, shape=shape) * 0.01\n",
        "\n",
        "    def three(rng):\n",
        "        return (normal(\n",
        "            (num_inputs, num_hiddens), rng), normal((num_hiddens, num_hiddens),\n",
        "                                                   rng), jnp.zeros(num_hiddens))\n",
        "\n",
        "    input_rng, forget_rng, output_rng, mem_rng, out_rng = jax.random.split(\n",
        "        init_rng, num=5)\n",
        "    W_xi, W_hi, b_i = three(input_rng)  # Input gate parameters\n",
        "    W_xf, W_hf, b_f = three(forget_rng)  # Forget gate parameters\n",
        "    W_xo, W_ho, b_o = three(output_rng)  # Output gate parameters\n",
        "    W_xc, W_hc, b_c = three(mem_rng)  # Candidate memory cell parameters\n",
        "    # Output layer parameters\n",
        "    W_hq = normal((num_hiddens, num_outputs), out_rng)\n",
        "    b_q = jnp.zeros(num_outputs)\n",
        "    params = [\n",
        "        W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c,\n",
        "        W_hq, b_q]\n",
        "    return params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "9yqm6PLhMBq9"
      },
      "outputs": [],
      "source": [
        "# The state is now a tuple of hidden state and cell state\n",
        "def init_lstm_state(batch_size, num_hiddens):\n",
        "    return (jnp.zeros((batch_size, num_hiddens)),\n",
        "            jnp.zeros((batch_size, num_hiddens)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "XB6S0k_MMJV-"
      },
      "outputs": [],
      "source": [
        "# forward function\n",
        "@jax.jit\n",
        "def lstm(params, state, inputs):\n",
        "    [W_xi, W_hi, b_i, W_xf, W_hf, b_f, W_xo, W_ho, b_o, W_xc, W_hc, b_c,\n",
        "        W_hq, b_q] = params\n",
        "    (H, C) = state\n",
        "    outputs = []\n",
        "    for X in inputs:\n",
        "        I = jax.nn.sigmoid((X @ W_xi) + (H @ W_hi) + b_i)\n",
        "        F = jax.nn.sigmoid((X @ W_xf) + (H @ W_hf) + b_f)\n",
        "        O = jax.nn.sigmoid((X @ W_xo) + (H @ W_ho) + b_o)\n",
        "        C_tilda = jnp.tanh((X @ W_xc) + (H @ W_hc) + b_c)\n",
        "        C = F * C + I * C_tilda\n",
        "        H = O * jnp.tanh(C)\n",
        "        Y = (H @ W_hq) + b_q\n",
        "        outputs.append(Y)\n",
        "    return jnp.concatenate(outputs, axis=0), (H, C)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YA59S3KlMQ0V"
      },
      "source": [
        "# Training and prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "Qw3vhoRKdVxt"
      },
      "outputs": [],
      "source": [
        "# Make the model class \n",
        "# Input X to call is (B,T) matrix of integers (from vocab encoding).\n",
        "# We transpose this to (T,B) then one-hot encode to (T,B,V), where V is vocab.\n",
        "# The result is passed to the forward function.\n",
        "# (We define the forward function as an argument, so we can change it later.)\n",
        "\n",
        "class RNNModelScratch: \n",
        "    \"\"\"A RNN Model implemented from scratch.\"\"\"\n",
        "    def __init__(self, vocab_size, num_hiddens, get_params, init_state,\n",
        "                 forward_fn):\n",
        "        self.vocab_size, self.num_hiddens = vocab_size, num_hiddens\n",
        "        self.init_state, self.get_params = init_state, get_params\n",
        "        self.forward_fn = forward_fn\n",
        "\n",
        "    def apply(self, params, state, X):\n",
        "        X = jax.nn.one_hot(X.T, num_classes=self.vocab_size)\n",
        "        return self.forward_fn(params, state, X)\n",
        "\n",
        "    def begin_state(self, batch_size):\n",
        "        return self.init_state(batch_size, self.num_hiddens)\n",
        "    \n",
        "    def init_params(self, rng):\n",
        "        return self.get_params(self.vocab_size, self.num_hiddens, rng)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "tA_QeZH5rp1o"
      },
      "outputs": [],
      "source": [
        "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",
        "        # Incrementally plot multiple lines\n",
        "        if legend is None:\n",
        "            legend = []\n",
        "        display.set_matplotlib_formats('svg')\n",
        "        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)\n",
        "        if nrows * ncols == 1:\n",
        "            self.axes = [self.axes,]\n",
        "        # Use a lambda function to capture arguments\n",
        "        self.config_axes = lambda: set_axes(self.axes[\n",
        "            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",
        "        # Add multiple data points into the figure\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 Timer:\n",
        "    \"\"\"Record multiple running times.\"\"\"\n",
        "    def __init__(self):\n",
        "        self.times = []\n",
        "        self.start()\n",
        "\n",
        "    def start(self):\n",
        "        \"\"\"Start the timer.\"\"\"\n",
        "        self.tik = time.time()\n",
        "\n",
        "    def stop(self):\n",
        "        \"\"\"Stop the timer and record the time in a list.\"\"\"\n",
        "        self.times.append(time.time() - self.tik)\n",
        "        return self.times[-1]\n",
        "\n",
        "    def avg(self):\n",
        "        \"\"\"Return the average time.\"\"\"\n",
        "        return sum(self.times) / len(self.times)\n",
        "\n",
        "    def sum(self):\n",
        "        \"\"\"Return the sum of time.\"\"\"\n",
        "        return sum(self.times)\n",
        "\n",
        "    def cumsum(self):\n",
        "        \"\"\"Return the accumulated time.\"\"\"\n",
        "        return jnp.array(self.times).cumsum().tolist()\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 set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):\n",
        "    \"\"\"Set the axes for matplotlib.\"\"\"\n",
        "    axes.set_xlabel(xlabel)\n",
        "    axes.set_ylabel(ylabel)\n",
        "    axes.set_xscale(xscale)\n",
        "    axes.set_yscale(yscale)\n",
        "    axes.set_xlim(xlim)\n",
        "    axes.set_ylim(ylim)\n",
        "    if legend:\n",
        "        axes.legend(legend)\n",
        "    axes.grid()\n",
        "\n",
        "@jax.jit\n",
        "def sgd(params, grads, lr, batch_size):\n",
        "    \"\"\"Minibatch stochastic gradient descent.\"\"\"\n",
        "    params = jax.tree_map(lambda p, g: p - lr * g / batch_size, params,\n",
        "                               grads)\n",
        "    return params\n",
        "\n",
        "@jax.jit\n",
        "def grad_clipping(grads, theta):\n",
        "    \"\"\"Clip the gradient.\"\"\"\n",
        "    def grad_update(grads):\n",
        "        return jax.tree_map(lambda g: g * theta / norm, grads)\n",
        "\n",
        "    norm = jnp.sqrt(sum(jax.tree_util.tree_leaves(jax.tree_map(\n",
        "        lambda x: jnp.sum(x ** 2), grads))))\n",
        "    # Update gradient if norm > theta\n",
        "    # This is jax.jit compatible\n",
        "    grads = jax.lax.cond(norm > theta, grad_update, lambda g: g, grads)\n",
        "    return grads"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "0tPGxQDWiqfl"
      },
      "outputs": [],
      "source": [
        "def predict(prefix, num_preds, net, params, vocab):  \n",
        "    \"\"\"Generate new characters following the `prefix`.\"\"\"\n",
        "    state = net.begin_state(batch_size=1)\n",
        "    outputs = [vocab[prefix[0]]]\n",
        "    get_input = lambda: jnp.array([outputs[-1]]).reshape((1, 1))\n",
        "    for y in prefix[1:]:  # Warm-up period\n",
        "        _, state = net.apply(params, state, get_input())\n",
        "        outputs.append(vocab[y])\n",
        "    for _ in range(num_preds):  # Predict `num_preds` steps\n",
        "        y, state = net.apply(params, state, get_input())\n",
        "        y = y.reshape(-1, y.shape[-1])\n",
        "        outputs.append(int(y.argmax(axis=1).reshape(1)))\n",
        "    return ''.join([vocab.idx_to_token[i] for i in outputs])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def train_step(apply_fn, loss_fn, params, state, X, Y):\n",
        "      def loss(params, state, X, Y):\n",
        "        y = Y.T.reshape(-1) #(B,T) -> (T,B)\n",
        "        y_hat, state = apply_fn(params, state, X)\n",
        "        y_hat = y_hat.reshape(-1, y_hat.shape[-1])\n",
        "        y_one_hot = jax.nn.one_hot(y, num_classes=y_hat.shape[-1])\n",
        "        return loss_fn(y_hat, y_one_hot).mean(), state\n",
        "\n",
        "      grad_fn = jax.value_and_grad(loss, has_aux=True)\n",
        "      (l, state), grads = grad_fn(params, state, X, Y)\n",
        "      grads = grad_clipping(grads, 1)\n",
        "      return l, state, grads"
      ],
      "metadata": {
        "id": "nAQ4RD0yh1vu"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_epoch(net, params, train_iter, loss, updater, use_random_iter):\n",
        "    state, timer = None, Timer()\n",
        "    metric = Accumulator(2)  # Sum of training loss, no. of tokens\n",
        "    if isinstance(updater, optax.GradientTransformation):\n",
        "      updater_state = updater.init(params)\n",
        "    # Convert to jax Partial functions for jax.jit compatibility \n",
        "    apply_fn = jax.tree_util.Partial(net.apply)\n",
        "    loss_fn = jax.tree_util.Partial(loss)\n",
        "    for X, Y in train_iter:\n",
        "        if state is None or use_random_iter:\n",
        "            # Initialize `state` when either it is the first iteration or\n",
        "            # using random sampling\n",
        "            state = net.begin_state(batch_size=X.shape[0])\n",
        "        l, state, grads = train_step(apply_fn, loss_fn, params, state, X, Y)\n",
        "        if isinstance(updater, optax.GradientTransformation):\n",
        "            updates, updater_state = updater.update(grads, updater_state)\n",
        "            params = optax.apply_updates(params, updates)\n",
        "        else:\n",
        "            # batch_size=1 since the `mean` function has been invoked\n",
        "            params = updater(params, grads, batch_size=1)\n",
        "        metric.add(l * Y.size, Y.size)\n",
        "    return params, math.exp(metric[0] / metric[1]), metric[1] / timer.stop()"
      ],
      "metadata": {
        "id": "GGkvH9UUh2mV"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(net, params, train_iter, vocab, lr, num_epochs,\n",
        "          use_random_iter=False):\n",
        "    loss = optax.softmax_cross_entropy\n",
        "    animator = Animator(xlabel='epoch', ylabel='perplexity',\n",
        "                            legend=['train'], xlim=[10, num_epochs])\n",
        "    # Initialize\n",
        "    if isinstance(net, nn.Module):\n",
        "        updater = optax.sgd(lr)\n",
        "    else:\n",
        "        updater = lambda params, grads, batch_size: sgd(params, grads, lr,\n",
        "                                                        batch_size)\n",
        "    num_preds = 50\n",
        "    predict_ = lambda prefix: predict(prefix, num_preds, net, params, vocab)\n",
        "    # Train and predict\n",
        "    for epoch in range(num_epochs):\n",
        "        params, ppl, speed = train_epoch(net, params, train_iter, loss, updater,\n",
        "                                         use_random_iter)\n",
        "        if (epoch + 1) % 10 == 0:\n",
        "            # Prediction takes time on the flax model\n",
        "            # print(predict_('time traveller'))\n",
        "            animator.add(epoch + 1, [ppl])\n",
        "    device = jax.default_backend()\n",
        "    print(f'perplexity {ppl:.1f}, {speed:.1f} tokens/sec on {device}')\n",
        "    print(predict_('time traveller'))\n",
        "    print(predict_('traveller'))\n",
        "    return params"
      ],
      "metadata": {
        "id": "epYm1kqHh4PW"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "random.seed(0)\n",
        "vocab_size, num_hiddens = len(vocab), 256\n",
        "num_epochs, lr = 500, 1\n",
        "model = RNNModelScratch(len(vocab), num_hiddens, get_lstm_params,\n",
        "                        init_lstm_state, lstm)\n",
        "params = model.init_params(rng)\n",
        "params = train(model, params, train_iter, vocab, lr, num_epochs)"
      ],
      "metadata": {
        "id": "8GO1pGfHiE87",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 314
        },
        "outputId": "0930806b-d3d0-40f4-81c2-66d55c21dad6"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "perplexity 1.1, 29941.3 tokens/sec on gpu\n",
            "time traveller for so it will be convenient to speak of himwas e\n",
            "travelleryou can show black is white by argument said filby\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 252x180 with 1 Axes>"
            ],
            "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Created with matplotlib (https://matplotlib.org/) -->\n<svg height=\"180.65625pt\" version=\"1.1\" viewBox=\"0 0 252.646875 180.65625\" width=\"252.646875pt\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n <defs>\n  <style type=\"text/css\">\n*{stroke-linecap:butt;stroke-linejoin:round;}\n  </style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 180.65625 \nL 252.646875 180.65625 \nL 252.646875 0 \nL 0 0 \nz\n\" style=\"fill:none;\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 40.603125 143.1 \nL 235.903125 143.1 \nL 235.903125 7.2 \nL 40.603125 7.2 \nz\n\" style=\"fill:#ffffff;\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 76.474554 143.1 \nL 76.474554 7.2 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path d=\"M 0 0 \nL 0 3.5 \n\" id=\"m98f689041d\" style=\"stroke:#000000;stroke-width:0.8;\"/>\n      </defs>\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"76.474554\" xlink:href=\"#m98f689041d\" y=\"143.1\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 100 -->\n      <defs>\n       <path d=\"M 12.40625 8.296875 \nL 28.515625 8.296875 \nL 28.515625 63.921875 \nL 10.984375 60.40625 \nL 10.984375 69.390625 \nL 28.421875 72.90625 \nL 38.28125 72.90625 \nL 38.28125 8.296875 \nL 54.390625 8.296875 \nL 54.390625 0 \nL 12.40625 0 \nz\n\" id=\"DejaVuSans-49\"/>\n       <path d=\"M 31.78125 66.40625 \nQ 24.171875 66.40625 20.328125 58.90625 \nQ 16.5 51.421875 16.5 36.375 \nQ 16.5 21.390625 20.328125 13.890625 \nQ 24.171875 6.390625 31.78125 6.390625 \nQ 39.453125 6.390625 43.28125 13.890625 \nQ 47.125 21.390625 47.125 36.375 \nQ 47.125 51.421875 43.28125 58.90625 \nQ 39.453125 66.40625 31.78125 66.40625 \nz\nM 31.78125 74.21875 \nQ 44.046875 74.21875 50.515625 64.515625 \nQ 56.984375 54.828125 56.984375 36.375 \nQ 56.984375 17.96875 50.515625 8.265625 \nQ 44.046875 -1.421875 31.78125 -1.421875 \nQ 19.53125 -1.421875 13.0625 8.265625 \nQ 6.59375 17.96875 6.59375 36.375 \nQ 6.59375 54.828125 13.0625 64.515625 \nQ 19.53125 74.21875 31.78125 74.21875 \nz\n\" id=\"DejaVuSans-48\"/>\n      </defs>\n      <g transform=\"translate(66.930804 157.698438)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 116.331696 143.1 \nL 116.331696 7.2 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"116.331696\" xlink:href=\"#m98f689041d\" y=\"143.1\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 200 -->\n      <defs>\n       <path d=\"M 19.1875 8.296875 \nL 53.609375 8.296875 \nL 53.609375 0 \nL 7.328125 0 \nL 7.328125 8.296875 \nQ 12.9375 14.109375 22.625 23.890625 \nQ 32.328125 33.6875 34.8125 36.53125 \nQ 39.546875 41.84375 41.421875 45.53125 \nQ 43.3125 49.21875 43.3125 52.78125 \nQ 43.3125 58.59375 39.234375 62.25 \nQ 35.15625 65.921875 28.609375 65.921875 \nQ 23.96875 65.921875 18.8125 64.3125 \nQ 13.671875 62.703125 7.8125 59.421875 \nL 7.8125 69.390625 \nQ 13.765625 71.78125 18.9375 73 \nQ 24.125 74.21875 28.421875 74.21875 \nQ 39.75 74.21875 46.484375 68.546875 \nQ 53.21875 62.890625 53.21875 53.421875 \nQ 53.21875 48.921875 51.53125 44.890625 \nQ 49.859375 40.875 45.40625 35.40625 \nQ 44.1875 33.984375 37.640625 27.21875 \nQ 31.109375 20.453125 19.1875 8.296875 \nz\n\" id=\"DejaVuSans-50\"/>\n      </defs>\n      <g transform=\"translate(106.787946 157.698438)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-50\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 156.188839 143.1 \nL 156.188839 7.2 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"156.188839\" xlink:href=\"#m98f689041d\" y=\"143.1\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 300 -->\n      <defs>\n       <path d=\"M 40.578125 39.3125 \nQ 47.65625 37.796875 51.625 33 \nQ 55.609375 28.21875 55.609375 21.1875 \nQ 55.609375 10.40625 48.1875 4.484375 \nQ 40.765625 -1.421875 27.09375 -1.421875 \nQ 22.515625 -1.421875 17.65625 -0.515625 \nQ 12.796875 0.390625 7.625 2.203125 \nL 7.625 11.71875 \nQ 11.71875 9.328125 16.59375 8.109375 \nQ 21.484375 6.890625 26.8125 6.890625 \nQ 36.078125 6.890625 40.9375 10.546875 \nQ 45.796875 14.203125 45.796875 21.1875 \nQ 45.796875 27.640625 41.28125 31.265625 \nQ 36.765625 34.90625 28.71875 34.90625 \nL 20.21875 34.90625 \nL 20.21875 43.015625 \nL 29.109375 43.015625 \nQ 36.375 43.015625 40.234375 45.921875 \nQ 44.09375 48.828125 44.09375 54.296875 \nQ 44.09375 59.90625 40.109375 62.90625 \nQ 36.140625 65.921875 28.71875 65.921875 \nQ 24.65625 65.921875 20.015625 65.03125 \nQ 15.375 64.15625 9.8125 62.3125 \nL 9.8125 71.09375 \nQ 15.4375 72.65625 20.34375 73.4375 \nQ 25.25 74.21875 29.59375 74.21875 \nQ 40.828125 74.21875 47.359375 69.109375 \nQ 53.90625 64.015625 53.90625 55.328125 \nQ 53.90625 49.265625 50.4375 45.09375 \nQ 46.96875 40.921875 40.578125 39.3125 \nz\n\" id=\"DejaVuSans-51\"/>\n      </defs>\n      <g transform=\"translate(146.645089 157.698438)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-51\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 196.045982 143.1 \nL 196.045982 7.2 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"196.045982\" xlink:href=\"#m98f689041d\" y=\"143.1\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 400 -->\n      <defs>\n       <path d=\"M 37.796875 64.3125 \nL 12.890625 25.390625 \nL 37.796875 25.390625 \nz\nM 35.203125 72.90625 \nL 47.609375 72.90625 \nL 47.609375 25.390625 \nL 58.015625 25.390625 \nL 58.015625 17.1875 \nL 47.609375 17.1875 \nL 47.609375 0 \nL 37.796875 0 \nL 37.796875 17.1875 \nL 4.890625 17.1875 \nL 4.890625 26.703125 \nz\n\" id=\"DejaVuSans-52\"/>\n      </defs>\n      <g transform=\"translate(186.502232 157.698438)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-52\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_5\">\n     <g id=\"line2d_9\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 235.903125 143.1 \nL 235.903125 7.2 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"235.903125\" xlink:href=\"#m98f689041d\" y=\"143.1\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 500 -->\n      <defs>\n       <path d=\"M 10.796875 72.90625 \nL 49.515625 72.90625 \nL 49.515625 64.59375 \nL 19.828125 64.59375 \nL 19.828125 46.734375 \nQ 21.96875 47.46875 24.109375 47.828125 \nQ 26.265625 48.1875 28.421875 48.1875 \nQ 40.625 48.1875 47.75 41.5 \nQ 54.890625 34.8125 54.890625 23.390625 \nQ 54.890625 11.625 47.5625 5.09375 \nQ 40.234375 -1.421875 26.90625 -1.421875 \nQ 22.3125 -1.421875 17.546875 -0.640625 \nQ 12.796875 0.140625 7.71875 1.703125 \nL 7.71875 11.625 \nQ 12.109375 9.234375 16.796875 8.0625 \nQ 21.484375 6.890625 26.703125 6.890625 \nQ 35.15625 6.890625 40.078125 11.328125 \nQ 45.015625 15.765625 45.015625 23.390625 \nQ 45.015625 31 40.078125 35.4375 \nQ 35.15625 39.890625 26.703125 39.890625 \nQ 22.75 39.890625 18.8125 39.015625 \nQ 14.890625 38.140625 10.796875 36.28125 \nz\n\" id=\"DejaVuSans-53\"/>\n      </defs>\n      <g transform=\"translate(226.359375 157.698438)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-53\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_6\">\n     <!-- epoch -->\n     <defs>\n      <path d=\"M 56.203125 29.59375 \nL 56.203125 25.203125 \nL 14.890625 25.203125 \nQ 15.484375 15.921875 20.484375 11.0625 \nQ 25.484375 6.203125 34.421875 6.203125 \nQ 39.59375 6.203125 44.453125 7.46875 \nQ 49.3125 8.734375 54.109375 11.28125 \nL 54.109375 2.78125 \nQ 49.265625 0.734375 44.1875 -0.34375 \nQ 39.109375 -1.421875 33.890625 -1.421875 \nQ 20.796875 -1.421875 13.15625 6.1875 \nQ 5.515625 13.8125 5.515625 26.8125 \nQ 5.515625 40.234375 12.765625 48.109375 \nQ 20.015625 56 32.328125 56 \nQ 43.359375 56 49.78125 48.890625 \nQ 56.203125 41.796875 56.203125 29.59375 \nz\nM 47.21875 32.234375 \nQ 47.125 39.59375 43.09375 43.984375 \nQ 39.0625 48.390625 32.421875 48.390625 \nQ 24.90625 48.390625 20.390625 44.140625 \nQ 15.875 39.890625 15.1875 32.171875 \nz\n\" id=\"DejaVuSans-101\"/>\n      <path d=\"M 18.109375 8.203125 \nL 18.109375 -20.796875 \nL 9.078125 -20.796875 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.390625 \nQ 20.953125 51.265625 25.265625 53.625 \nQ 29.59375 56 35.59375 56 \nQ 45.5625 56 51.78125 48.09375 \nQ 58.015625 40.1875 58.015625 27.296875 \nQ 58.015625 14.40625 51.78125 6.484375 \nQ 45.5625 -1.421875 35.59375 -1.421875 \nQ 29.59375 -1.421875 25.265625 0.953125 \nQ 20.953125 3.328125 18.109375 8.203125 \nz\nM 48.6875 27.296875 \nQ 48.6875 37.203125 44.609375 42.84375 \nQ 40.53125 48.484375 33.40625 48.484375 \nQ 26.265625 48.484375 22.1875 42.84375 \nQ 18.109375 37.203125 18.109375 27.296875 \nQ 18.109375 17.390625 22.1875 11.75 \nQ 26.265625 6.109375 33.40625 6.109375 \nQ 40.53125 6.109375 44.609375 11.75 \nQ 48.6875 17.390625 48.6875 27.296875 \nz\n\" id=\"DejaVuSans-112\"/>\n      <path d=\"M 30.609375 48.390625 \nQ 23.390625 48.390625 19.1875 42.75 \nQ 14.984375 37.109375 14.984375 27.296875 \nQ 14.984375 17.484375 19.15625 11.84375 \nQ 23.34375 6.203125 30.609375 6.203125 \nQ 37.796875 6.203125 41.984375 11.859375 \nQ 46.1875 17.53125 46.1875 27.296875 \nQ 46.1875 37.015625 41.984375 42.703125 \nQ 37.796875 48.390625 30.609375 48.390625 \nz\nM 30.609375 56 \nQ 42.328125 56 49.015625 48.375 \nQ 55.71875 40.765625 55.71875 27.296875 \nQ 55.71875 13.875 49.015625 6.21875 \nQ 42.328125 -1.421875 30.609375 -1.421875 \nQ 18.84375 -1.421875 12.171875 6.21875 \nQ 5.515625 13.875 5.515625 27.296875 \nQ 5.515625 40.765625 12.171875 48.375 \nQ 18.84375 56 30.609375 56 \nz\n\" id=\"DejaVuSans-111\"/>\n      <path d=\"M 48.78125 52.59375 \nL 48.78125 44.1875 \nQ 44.96875 46.296875 41.140625 47.34375 \nQ 37.3125 48.390625 33.40625 48.390625 \nQ 24.65625 48.390625 19.8125 42.84375 \nQ 14.984375 37.3125 14.984375 27.296875 \nQ 14.984375 17.28125 19.8125 11.734375 \nQ 24.65625 6.203125 33.40625 6.203125 \nQ 37.3125 6.203125 41.140625 7.25 \nQ 44.96875 8.296875 48.78125 10.40625 \nL 48.78125 2.09375 \nQ 45.015625 0.34375 40.984375 -0.53125 \nQ 36.96875 -1.421875 32.421875 -1.421875 \nQ 20.0625 -1.421875 12.78125 6.34375 \nQ 5.515625 14.109375 5.515625 27.296875 \nQ 5.515625 40.671875 12.859375 48.328125 \nQ 20.21875 56 33.015625 56 \nQ 37.15625 56 41.109375 55.140625 \nQ 45.0625 54.296875 48.78125 52.59375 \nz\n\" id=\"DejaVuSans-99\"/>\n      <path d=\"M 54.890625 33.015625 \nL 54.890625 0 \nL 45.90625 0 \nL 45.90625 32.71875 \nQ 45.90625 40.484375 42.875 44.328125 \nQ 39.84375 48.1875 33.796875 48.1875 \nQ 26.515625 48.1875 22.3125 43.546875 \nQ 18.109375 38.921875 18.109375 30.90625 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 75.984375 \nL 18.109375 75.984375 \nL 18.109375 46.1875 \nQ 21.34375 51.125 25.703125 53.5625 \nQ 30.078125 56 35.796875 56 \nQ 45.21875 56 50.046875 50.171875 \nQ 54.890625 44.34375 54.890625 33.015625 \nz\n\" id=\"DejaVuSans-104\"/>\n     </defs>\n     <g transform=\"translate(123.025 171.376563)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"61.523438\" xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"125\" xlink:href=\"#DejaVuSans-111\"/>\n      <use x=\"186.181641\" xlink:href=\"#DejaVuSans-99\"/>\n      <use x=\"241.162109\" xlink:href=\"#DejaVuSans-104\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_11\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 40.603125 108.139365 \nL 235.903125 108.139365 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <defs>\n       <path d=\"M 0 0 \nL -3.5 0 \n\" id=\"m9cfc341d0a\" style=\"stroke:#000000;stroke-width:0.8;\"/>\n      </defs>\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"40.603125\" xlink:href=\"#m9cfc341d0a\" y=\"108.139365\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 5 -->\n      <g transform=\"translate(27.240625 111.938584)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-53\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_13\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 40.603125 71.385076 \nL 235.903125 71.385076 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"40.603125\" xlink:href=\"#m9cfc341d0a\" y=\"71.385076\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 10 -->\n      <g transform=\"translate(20.878125 75.184295)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_15\">\n      <path clip-path=\"url(#p8d53356726)\" d=\"M 40.603125 34.630787 \nL 235.903125 34.630787 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"40.603125\" xlink:href=\"#m9cfc341d0a\" y=\"34.630787\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 15 -->\n      <g transform=\"translate(20.878125 38.430006)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-53\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_10\">\n     <!-- perplexity -->\n     <defs>\n      <path d=\"M 41.109375 46.296875 \nQ 39.59375 47.171875 37.8125 47.578125 \nQ 36.03125 48 33.890625 48 \nQ 26.265625 48 22.1875 43.046875 \nQ 18.109375 38.09375 18.109375 28.8125 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 20.953125 51.171875 25.484375 53.578125 \nQ 30.03125 56 36.53125 56 \nQ 37.453125 56 38.578125 55.875 \nQ 39.703125 55.765625 41.0625 55.515625 \nz\n\" id=\"DejaVuSans-114\"/>\n      <path d=\"M 9.421875 75.984375 \nL 18.40625 75.984375 \nL 18.40625 0 \nL 9.421875 0 \nz\n\" id=\"DejaVuSans-108\"/>\n      <path d=\"M 54.890625 54.6875 \nL 35.109375 28.078125 \nL 55.90625 0 \nL 45.3125 0 \nL 29.390625 21.484375 \nL 13.484375 0 \nL 2.875 0 \nL 24.125 28.609375 \nL 4.6875 54.6875 \nL 15.28125 54.6875 \nL 29.78125 35.203125 \nL 44.28125 54.6875 \nz\n\" id=\"DejaVuSans-120\"/>\n      <path d=\"M 9.421875 54.6875 \nL 18.40625 54.6875 \nL 18.40625 0 \nL 9.421875 0 \nz\nM 9.421875 75.984375 \nL 18.40625 75.984375 \nL 18.40625 64.59375 \nL 9.421875 64.59375 \nz\n\" id=\"DejaVuSans-105\"/>\n      <path d=\"M 18.3125 70.21875 \nL 18.3125 54.6875 \nL 36.8125 54.6875 \nL 36.8125 47.703125 \nL 18.3125 47.703125 \nL 18.3125 18.015625 \nQ 18.3125 11.328125 20.140625 9.421875 \nQ 21.96875 7.515625 27.59375 7.515625 \nL 36.8125 7.515625 \nL 36.8125 0 \nL 27.59375 0 \nQ 17.1875 0 13.234375 3.875 \nQ 9.28125 7.765625 9.28125 18.015625 \nL 9.28125 47.703125 \nL 2.6875 47.703125 \nL 2.6875 54.6875 \nL 9.28125 54.6875 \nL 9.28125 70.21875 \nz\n\" id=\"DejaVuSans-116\"/>\n      <path d=\"M 32.171875 -5.078125 \nQ 28.375 -14.84375 24.75 -17.8125 \nQ 21.140625 -20.796875 15.09375 -20.796875 \nL 7.90625 -20.796875 \nL 7.90625 -13.28125 \nL 13.1875 -13.28125 \nQ 16.890625 -13.28125 18.9375 -11.515625 \nQ 21 -9.765625 23.484375 -3.21875 \nL 25.09375 0.875 \nL 2.984375 54.6875 \nL 12.5 54.6875 \nL 29.59375 11.921875 \nL 46.6875 54.6875 \nL 56.203125 54.6875 \nz\n\" id=\"DejaVuSans-121\"/>\n     </defs>\n     <g transform=\"translate(14.798437 100.276562)rotate(-90)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"63.476562\" xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"125\" xlink:href=\"#DejaVuSans-114\"/>\n      <use x=\"166.113281\" xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"229.589844\" xlink:href=\"#DejaVuSans-108\"/>\n      <use x=\"257.373047\" xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"317.146484\" xlink:href=\"#DejaVuSans-120\"/>\n      <use x=\"376.326172\" xlink:href=\"#DejaVuSans-105\"/>\n      <use x=\"404.109375\" xlink:href=\"#DejaVuSans-116\"/>\n      <use x=\"443.318359\" xlink:href=\"#DejaVuSans-121\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_17\">\n    <path clip-path=\"url(#p8d53356726)\" d=\"M 40.603125 13.377273 \nL 44.588839 16.71519 \nL 48.574554 22.096155 \nL 52.560268 29.549881 \nL 56.545982 37.766725 \nL 60.531696 49.543741 \nL 64.517411 57.781888 \nL 68.503125 61.726239 \nL 72.488839 64.704262 \nL 76.474554 67.513489 \nL 80.460268 69.972964 \nL 84.445982 71.274342 \nL 88.431696 74.182644 \nL 92.417411 76.512587 \nL 96.403125 78.56349 \nL 100.388839 80.874669 \nL 104.374554 82.882116 \nL 108.360268 85.893985 \nL 112.345982 88.378876 \nL 116.331696 90.038767 \nL 120.317411 92.016135 \nL 124.303125 94.017665 \nL 128.288839 95.318327 \nL 132.274554 99.179341 \nL 136.260268 101.080204 \nL 140.245982 103.63522 \nL 144.231696 106.711339 \nL 148.217411 108.353653 \nL 152.203125 111.676993 \nL 156.188839 113.955708 \nL 160.174554 116.028268 \nL 164.160268 118.224033 \nL 168.145982 120.92416 \nL 172.131696 123.49765 \nL 176.117411 125.466036 \nL 180.103125 126.722447 \nL 184.088839 129.233121 \nL 188.074554 130.978849 \nL 192.060268 131.70725 \nL 196.045982 133.117158 \nL 200.031696 133.766877 \nL 204.017411 134.556967 \nL 208.003125 134.378206 \nL 211.988839 135.087358 \nL 215.974554 135.985937 \nL 219.960268 136.441004 \nL 223.945982 136.649289 \nL 227.931696 135.921332 \nL 231.917411 136.879517 \nL 235.903125 136.922727 \n\" style=\"fill:none;stroke:#1f77b4;stroke-linecap:square;stroke-width:1.5;\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 40.603125 143.1 \nL 40.603125 7.2 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 235.903125 143.1 \nL 235.903125 7.2 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 40.603125 143.1 \nL 235.903125 143.1 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 40.603125 7.2 \nL 235.903125 7.2 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 173.628125 29.878125 \nL 228.903125 29.878125 \nQ 230.903125 29.878125 230.903125 27.878125 \nL 230.903125 14.2 \nQ 230.903125 12.2 228.903125 12.2 \nL 173.628125 12.2 \nQ 171.628125 12.2 171.628125 14.2 \nL 171.628125 27.878125 \nQ 171.628125 29.878125 173.628125 29.878125 \nz\n\" style=\"fill:#ffffff;opacity:0.8;stroke:#cccccc;stroke-linejoin:miter;\"/>\n    </g>\n    <g id=\"line2d_18\">\n     <path d=\"M 175.628125 20.298437 \nL 195.628125 20.298437 \n\" style=\"fill:none;stroke:#1f77b4;stroke-linecap:square;stroke-width:1.5;\"/>\n    </g>\n    <g id=\"line2d_19\"/>\n    <g id=\"text_11\">\n     <!-- train -->\n     <defs>\n      <path d=\"M 34.28125 27.484375 \nQ 23.390625 27.484375 19.1875 25 \nQ 14.984375 22.515625 14.984375 16.5 \nQ 14.984375 11.71875 18.140625 8.90625 \nQ 21.296875 6.109375 26.703125 6.109375 \nQ 34.1875 6.109375 38.703125 11.40625 \nQ 43.21875 16.703125 43.21875 25.484375 \nL 43.21875 27.484375 \nz\nM 52.203125 31.203125 \nL 52.203125 0 \nL 43.21875 0 \nL 43.21875 8.296875 \nQ 40.140625 3.328125 35.546875 0.953125 \nQ 30.953125 -1.421875 24.3125 -1.421875 \nQ 15.921875 -1.421875 10.953125 3.296875 \nQ 6 8.015625 6 15.921875 \nQ 6 25.140625 12.171875 29.828125 \nQ 18.359375 34.515625 30.609375 34.515625 \nL 43.21875 34.515625 \nL 43.21875 35.40625 \nQ 43.21875 41.609375 39.140625 45 \nQ 35.0625 48.390625 27.6875 48.390625 \nQ 23 48.390625 18.546875 47.265625 \nQ 14.109375 46.140625 10.015625 43.890625 \nL 10.015625 52.203125 \nQ 14.9375 54.109375 19.578125 55.046875 \nQ 24.21875 56 28.609375 56 \nQ 40.484375 56 46.34375 49.84375 \nQ 52.203125 43.703125 52.203125 31.203125 \nz\n\" id=\"DejaVuSans-97\"/>\n      <path d=\"M 54.890625 33.015625 \nL 54.890625 0 \nL 45.90625 0 \nL 45.90625 32.71875 \nQ 45.90625 40.484375 42.875 44.328125 \nQ 39.84375 48.1875 33.796875 48.1875 \nQ 26.515625 48.1875 22.3125 43.546875 \nQ 18.109375 38.921875 18.109375 30.90625 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 21.34375 51.125 25.703125 53.5625 \nQ 30.078125 56 35.796875 56 \nQ 45.21875 56 50.046875 50.171875 \nQ 54.890625 44.34375 54.890625 33.015625 \nz\n\" id=\"DejaVuSans-110\"/>\n     </defs>\n     <g transform=\"translate(203.628125 23.798437)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-116\"/>\n      <use x=\"39.208984\" xlink:href=\"#DejaVuSans-114\"/>\n      <use x=\"80.322266\" xlink:href=\"#DejaVuSans-97\"/>\n      <use x=\"141.601562\" xlink:href=\"#DejaVuSans-105\"/>\n      <use x=\"169.384766\" xlink:href=\"#DejaVuSans-110\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p8d53356726\">\n   <rect height=\"135.9\" width=\"195.3\" x=\"40.603125\" y=\"7.2\"/>\n  </clipPath>\n </defs>\n</svg>\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GkwrwyP_MVif"
      },
      "source": [
        "## Using flax module"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class LSTM(nn.Module):\n",
        "    \"\"\"The LSTM layer.\"\"\"\n",
        "    @functools.partial(\n",
        "        nn.transforms.scan,\n",
        "        variable_broadcast='params',\n",
        "        in_axes=0, out_axes=0,\n",
        "        split_rngs={'params': False})\n",
        "    @nn.compact\n",
        "    def __call__(self, state, x):\n",
        "        return nn.OptimizedLSTMCell()(state, x)\n",
        "\n",
        "    @staticmethod\n",
        "    def initialize_carry(rng, batch_dims, size):\n",
        "        return nn.OptimizedLSTMCell.initialize_carry(rng, batch_dims, size)"
      ],
      "metadata": {
        "id": "yNRVFTN_iVbW"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "BnfRIJXn9Irv"
      },
      "outputs": [],
      "source": [
        "class RNNModel(nn.Module):\n",
        "    \"\"\"The RNN model.\"\"\"\n",
        "    rnn: nn.Module\n",
        "    vocab_size: int\n",
        "    num_hiddens: int\n",
        "    bidirectional: bool = False\n",
        "        \n",
        "    def setup(self):\n",
        "        # If the RNN is bidirectional (to be introduced later),\n",
        "        # `num_directions` should be 2, else it should be 1.\n",
        "        if not self.bidirectional:\n",
        "            self.num_directions = 1\n",
        "        else:\n",
        "            self.num_directions = 2\n",
        "\n",
        "    @nn.compact\n",
        "    def __call__(self, state, inputs):\n",
        "        X = jax.nn.one_hot(inputs.T, num_classes=self.vocab_size)\n",
        "        state, Y = self.rnn(state, X)\n",
        "        output = nn.Dense(self.vocab_size)(Y)\n",
        "        return output, state\n",
        "\n",
        "    def begin_state(self, batch_size=1):\n",
        "        # Use fixed random key since default state init fn is just `zeros`.\n",
        "        return self.rnn.initialize_carry(jax.random.PRNGKey(0), (batch_size,),\n",
        "                                         self.num_hiddens)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "random.seed(0)\n",
        "lstm_layer = LSTM()\n",
        "model = RNNModel(rnn=lstm_layer, vocab_size=len(vocab), num_hiddens=num_hiddens)\n",
        "initial_state = model.begin_state(batch_size)\n",
        "params = model.init(rng, initial_state, jnp.ones([batch_size, num_steps]))\n",
        "params = train(model, params, train_iter, vocab, lr, num_epochs)"
      ],
      "metadata": {
        "id": "FK8EePtSjAEh",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 318
        },
        "outputId": "a151c55b-0795-41a2-d3e8-d521819cb45c"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "perplexity 1.0, 24683.2 tokens/sec on gpu\n",
            "time travelleryou can show black is white by argument said filby\n",
            "travelleryou can show black is white by argument said filby\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 252x180 with 1 Axes>"
            ],
            "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Created with matplotlib (https://matplotlib.org/) -->\n<svg height=\"183.791071pt\" version=\"1.1\" viewBox=\"0 0 262.1875 183.791071\" width=\"262.1875pt\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n <defs>\n  <style type=\"text/css\">\n*{stroke-linecap:butt;stroke-linejoin:round;}\n  </style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.791071 \nL 262.1875 183.791071 \nL 262.1875 0 \nL 0 0 \nz\n\" style=\"fill:none;\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 50.14375 146.234821 \nL 245.44375 146.234821 \nL 245.44375 10.334821 \nL 50.14375 10.334821 \nz\n\" style=\"fill:#ffffff;\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 86.015179 146.234821 \nL 86.015179 10.334821 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path d=\"M 0 0 \nL 0 3.5 \n\" id=\"m3e8e10fd78\" style=\"stroke:#000000;stroke-width:0.8;\"/>\n      </defs>\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"86.015179\" xlink:href=\"#m3e8e10fd78\" y=\"146.234821\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 100 -->\n      <defs>\n       <path d=\"M 12.40625 8.296875 \nL 28.515625 8.296875 \nL 28.515625 63.921875 \nL 10.984375 60.40625 \nL 10.984375 69.390625 \nL 28.421875 72.90625 \nL 38.28125 72.90625 \nL 38.28125 8.296875 \nL 54.390625 8.296875 \nL 54.390625 0 \nL 12.40625 0 \nz\n\" id=\"DejaVuSans-49\"/>\n       <path d=\"M 31.78125 66.40625 \nQ 24.171875 66.40625 20.328125 58.90625 \nQ 16.5 51.421875 16.5 36.375 \nQ 16.5 21.390625 20.328125 13.890625 \nQ 24.171875 6.390625 31.78125 6.390625 \nQ 39.453125 6.390625 43.28125 13.890625 \nQ 47.125 21.390625 47.125 36.375 \nQ 47.125 51.421875 43.28125 58.90625 \nQ 39.453125 66.40625 31.78125 66.40625 \nz\nM 31.78125 74.21875 \nQ 44.046875 74.21875 50.515625 64.515625 \nQ 56.984375 54.828125 56.984375 36.375 \nQ 56.984375 17.96875 50.515625 8.265625 \nQ 44.046875 -1.421875 31.78125 -1.421875 \nQ 19.53125 -1.421875 13.0625 8.265625 \nQ 6.59375 17.96875 6.59375 36.375 \nQ 6.59375 54.828125 13.0625 64.515625 \nQ 19.53125 74.21875 31.78125 74.21875 \nz\n\" id=\"DejaVuSans-48\"/>\n      </defs>\n      <g transform=\"translate(76.471429 160.833258)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 125.872321 146.234821 \nL 125.872321 10.334821 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"125.872321\" xlink:href=\"#m3e8e10fd78\" y=\"146.234821\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 200 -->\n      <defs>\n       <path d=\"M 19.1875 8.296875 \nL 53.609375 8.296875 \nL 53.609375 0 \nL 7.328125 0 \nL 7.328125 8.296875 \nQ 12.9375 14.109375 22.625 23.890625 \nQ 32.328125 33.6875 34.8125 36.53125 \nQ 39.546875 41.84375 41.421875 45.53125 \nQ 43.3125 49.21875 43.3125 52.78125 \nQ 43.3125 58.59375 39.234375 62.25 \nQ 35.15625 65.921875 28.609375 65.921875 \nQ 23.96875 65.921875 18.8125 64.3125 \nQ 13.671875 62.703125 7.8125 59.421875 \nL 7.8125 69.390625 \nQ 13.765625 71.78125 18.9375 73 \nQ 24.125 74.21875 28.421875 74.21875 \nQ 39.75 74.21875 46.484375 68.546875 \nQ 53.21875 62.890625 53.21875 53.421875 \nQ 53.21875 48.921875 51.53125 44.890625 \nQ 49.859375 40.875 45.40625 35.40625 \nQ 44.1875 33.984375 37.640625 27.21875 \nQ 31.109375 20.453125 19.1875 8.296875 \nz\n\" id=\"DejaVuSans-50\"/>\n      </defs>\n      <g transform=\"translate(116.328571 160.833258)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-50\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 165.729464 146.234821 \nL 165.729464 10.334821 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"165.729464\" xlink:href=\"#m3e8e10fd78\" y=\"146.234821\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 300 -->\n      <defs>\n       <path d=\"M 40.578125 39.3125 \nQ 47.65625 37.796875 51.625 33 \nQ 55.609375 28.21875 55.609375 21.1875 \nQ 55.609375 10.40625 48.1875 4.484375 \nQ 40.765625 -1.421875 27.09375 -1.421875 \nQ 22.515625 -1.421875 17.65625 -0.515625 \nQ 12.796875 0.390625 7.625 2.203125 \nL 7.625 11.71875 \nQ 11.71875 9.328125 16.59375 8.109375 \nQ 21.484375 6.890625 26.8125 6.890625 \nQ 36.078125 6.890625 40.9375 10.546875 \nQ 45.796875 14.203125 45.796875 21.1875 \nQ 45.796875 27.640625 41.28125 31.265625 \nQ 36.765625 34.90625 28.71875 34.90625 \nL 20.21875 34.90625 \nL 20.21875 43.015625 \nL 29.109375 43.015625 \nQ 36.375 43.015625 40.234375 45.921875 \nQ 44.09375 48.828125 44.09375 54.296875 \nQ 44.09375 59.90625 40.109375 62.90625 \nQ 36.140625 65.921875 28.71875 65.921875 \nQ 24.65625 65.921875 20.015625 65.03125 \nQ 15.375 64.15625 9.8125 62.3125 \nL 9.8125 71.09375 \nQ 15.4375 72.65625 20.34375 73.4375 \nQ 25.25 74.21875 29.59375 74.21875 \nQ 40.828125 74.21875 47.359375 69.109375 \nQ 53.90625 64.015625 53.90625 55.328125 \nQ 53.90625 49.265625 50.4375 45.09375 \nQ 46.96875 40.921875 40.578125 39.3125 \nz\n\" id=\"DejaVuSans-51\"/>\n      </defs>\n      <g transform=\"translate(156.185714 160.833258)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-51\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 205.586607 146.234821 \nL 205.586607 10.334821 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"205.586607\" xlink:href=\"#m3e8e10fd78\" y=\"146.234821\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 400 -->\n      <defs>\n       <path d=\"M 37.796875 64.3125 \nL 12.890625 25.390625 \nL 37.796875 25.390625 \nz\nM 35.203125 72.90625 \nL 47.609375 72.90625 \nL 47.609375 25.390625 \nL 58.015625 25.390625 \nL 58.015625 17.1875 \nL 47.609375 17.1875 \nL 47.609375 0 \nL 37.796875 0 \nL 37.796875 17.1875 \nL 4.890625 17.1875 \nL 4.890625 26.703125 \nz\n\" id=\"DejaVuSans-52\"/>\n      </defs>\n      <g transform=\"translate(196.042857 160.833258)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-52\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_5\">\n     <g id=\"line2d_9\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 245.44375 146.234821 \nL 245.44375 10.334821 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"245.44375\" xlink:href=\"#m3e8e10fd78\" y=\"146.234821\"/>\n      </g>\n     </g>\n     <g id=\"text_5\">\n      <!-- 500 -->\n      <defs>\n       <path d=\"M 10.796875 72.90625 \nL 49.515625 72.90625 \nL 49.515625 64.59375 \nL 19.828125 64.59375 \nL 19.828125 46.734375 \nQ 21.96875 47.46875 24.109375 47.828125 \nQ 26.265625 48.1875 28.421875 48.1875 \nQ 40.625 48.1875 47.75 41.5 \nQ 54.890625 34.8125 54.890625 23.390625 \nQ 54.890625 11.625 47.5625 5.09375 \nQ 40.234375 -1.421875 26.90625 -1.421875 \nQ 22.3125 -1.421875 17.546875 -0.640625 \nQ 12.796875 0.140625 7.71875 1.703125 \nL 7.71875 11.625 \nQ 12.109375 9.234375 16.796875 8.0625 \nQ 21.484375 6.890625 26.703125 6.890625 \nQ 35.15625 6.890625 40.078125 11.328125 \nQ 45.015625 15.765625 45.015625 23.390625 \nQ 45.015625 31 40.078125 35.4375 \nQ 35.15625 39.890625 26.703125 39.890625 \nQ 22.75 39.890625 18.8125 39.015625 \nQ 14.890625 38.140625 10.796875 36.28125 \nz\n\" id=\"DejaVuSans-53\"/>\n      </defs>\n      <g transform=\"translate(235.9 160.833258)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-53\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_6\">\n     <!-- epoch -->\n     <defs>\n      <path d=\"M 56.203125 29.59375 \nL 56.203125 25.203125 \nL 14.890625 25.203125 \nQ 15.484375 15.921875 20.484375 11.0625 \nQ 25.484375 6.203125 34.421875 6.203125 \nQ 39.59375 6.203125 44.453125 7.46875 \nQ 49.3125 8.734375 54.109375 11.28125 \nL 54.109375 2.78125 \nQ 49.265625 0.734375 44.1875 -0.34375 \nQ 39.109375 -1.421875 33.890625 -1.421875 \nQ 20.796875 -1.421875 13.15625 6.1875 \nQ 5.515625 13.8125 5.515625 26.8125 \nQ 5.515625 40.234375 12.765625 48.109375 \nQ 20.015625 56 32.328125 56 \nQ 43.359375 56 49.78125 48.890625 \nQ 56.203125 41.796875 56.203125 29.59375 \nz\nM 47.21875 32.234375 \nQ 47.125 39.59375 43.09375 43.984375 \nQ 39.0625 48.390625 32.421875 48.390625 \nQ 24.90625 48.390625 20.390625 44.140625 \nQ 15.875 39.890625 15.1875 32.171875 \nz\n\" id=\"DejaVuSans-101\"/>\n      <path d=\"M 18.109375 8.203125 \nL 18.109375 -20.796875 \nL 9.078125 -20.796875 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.390625 \nQ 20.953125 51.265625 25.265625 53.625 \nQ 29.59375 56 35.59375 56 \nQ 45.5625 56 51.78125 48.09375 \nQ 58.015625 40.1875 58.015625 27.296875 \nQ 58.015625 14.40625 51.78125 6.484375 \nQ 45.5625 -1.421875 35.59375 -1.421875 \nQ 29.59375 -1.421875 25.265625 0.953125 \nQ 20.953125 3.328125 18.109375 8.203125 \nz\nM 48.6875 27.296875 \nQ 48.6875 37.203125 44.609375 42.84375 \nQ 40.53125 48.484375 33.40625 48.484375 \nQ 26.265625 48.484375 22.1875 42.84375 \nQ 18.109375 37.203125 18.109375 27.296875 \nQ 18.109375 17.390625 22.1875 11.75 \nQ 26.265625 6.109375 33.40625 6.109375 \nQ 40.53125 6.109375 44.609375 11.75 \nQ 48.6875 17.390625 48.6875 27.296875 \nz\n\" id=\"DejaVuSans-112\"/>\n      <path d=\"M 30.609375 48.390625 \nQ 23.390625 48.390625 19.1875 42.75 \nQ 14.984375 37.109375 14.984375 27.296875 \nQ 14.984375 17.484375 19.15625 11.84375 \nQ 23.34375 6.203125 30.609375 6.203125 \nQ 37.796875 6.203125 41.984375 11.859375 \nQ 46.1875 17.53125 46.1875 27.296875 \nQ 46.1875 37.015625 41.984375 42.703125 \nQ 37.796875 48.390625 30.609375 48.390625 \nz\nM 30.609375 56 \nQ 42.328125 56 49.015625 48.375 \nQ 55.71875 40.765625 55.71875 27.296875 \nQ 55.71875 13.875 49.015625 6.21875 \nQ 42.328125 -1.421875 30.609375 -1.421875 \nQ 18.84375 -1.421875 12.171875 6.21875 \nQ 5.515625 13.875 5.515625 27.296875 \nQ 5.515625 40.765625 12.171875 48.375 \nQ 18.84375 56 30.609375 56 \nz\n\" id=\"DejaVuSans-111\"/>\n      <path d=\"M 48.78125 52.59375 \nL 48.78125 44.1875 \nQ 44.96875 46.296875 41.140625 47.34375 \nQ 37.3125 48.390625 33.40625 48.390625 \nQ 24.65625 48.390625 19.8125 42.84375 \nQ 14.984375 37.3125 14.984375 27.296875 \nQ 14.984375 17.28125 19.8125 11.734375 \nQ 24.65625 6.203125 33.40625 6.203125 \nQ 37.3125 6.203125 41.140625 7.25 \nQ 44.96875 8.296875 48.78125 10.40625 \nL 48.78125 2.09375 \nQ 45.015625 0.34375 40.984375 -0.53125 \nQ 36.96875 -1.421875 32.421875 -1.421875 \nQ 20.0625 -1.421875 12.78125 6.34375 \nQ 5.515625 14.109375 5.515625 27.296875 \nQ 5.515625 40.671875 12.859375 48.328125 \nQ 20.21875 56 33.015625 56 \nQ 37.15625 56 41.109375 55.140625 \nQ 45.0625 54.296875 48.78125 52.59375 \nz\n\" id=\"DejaVuSans-99\"/>\n      <path d=\"M 54.890625 33.015625 \nL 54.890625 0 \nL 45.90625 0 \nL 45.90625 32.71875 \nQ 45.90625 40.484375 42.875 44.328125 \nQ 39.84375 48.1875 33.796875 48.1875 \nQ 26.515625 48.1875 22.3125 43.546875 \nQ 18.109375 38.921875 18.109375 30.90625 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 75.984375 \nL 18.109375 75.984375 \nL 18.109375 46.1875 \nQ 21.34375 51.125 25.703125 53.5625 \nQ 30.078125 56 35.796875 56 \nQ 45.21875 56 50.046875 50.171875 \nQ 54.890625 44.34375 54.890625 33.015625 \nz\n\" id=\"DejaVuSans-104\"/>\n     </defs>\n     <g transform=\"translate(132.565625 174.511383)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"61.523438\" xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"125\" xlink:href=\"#DejaVuSans-111\"/>\n      <use x=\"186.181641\" xlink:href=\"#DejaVuSans-99\"/>\n      <use x=\"241.162109\" xlink:href=\"#DejaVuSans-104\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_11\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 126.49777 \nL 245.44375 126.49777 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <defs>\n       <path d=\"M 0 0 \nL -3.5 0 \n\" id=\"m389b38bd27\" style=\"stroke:#000000;stroke-width:0.8;\"/>\n      </defs>\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"126.49777\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 2.5 -->\n      <defs>\n       <path d=\"M 10.6875 12.40625 \nL 21 12.40625 \nL 21 0 \nL 10.6875 0 \nz\n\" id=\"DejaVuSans-46\"/>\n      </defs>\n      <g transform=\"translate(27.240625 130.296988)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-50\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"95.410156\" xlink:href=\"#DejaVuSans-53\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_13\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 103.398059 \nL 245.44375 103.398059 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"103.398059\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 5.0 -->\n      <g transform=\"translate(27.240625 107.197278)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-53\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"95.410156\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_15\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 80.298349 \nL 245.44375 80.298349 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"80.298349\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 7.5 -->\n      <defs>\n       <path d=\"M 8.203125 72.90625 \nL 55.078125 72.90625 \nL 55.078125 68.703125 \nL 28.609375 0 \nL 18.3125 0 \nL 43.21875 64.59375 \nL 8.203125 64.59375 \nz\n\" id=\"DejaVuSans-55\"/>\n      </defs>\n      <g transform=\"translate(27.240625 84.097568)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-55\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"95.410156\" xlink:href=\"#DejaVuSans-53\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_17\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 57.198639 \nL 245.44375 57.198639 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_18\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"57.198639\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 10.0 -->\n      <g transform=\"translate(20.878125 60.997858)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-48\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"159.033203\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_19\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 34.098929 \nL 245.44375 34.098929 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_20\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"34.098929\"/>\n      </g>\n     </g>\n     <g id=\"text_11\">\n      <!-- 12.5 -->\n      <g transform=\"translate(20.878125 37.898148)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-50\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"159.033203\" xlink:href=\"#DejaVuSans-53\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_6\">\n     <g id=\"line2d_21\">\n      <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 10.999219 \nL 245.44375 10.999219 \n\" style=\"fill:none;stroke:#b0b0b0;stroke-linecap:square;stroke-width:0.8;\"/>\n     </g>\n     <g id=\"line2d_22\">\n      <g>\n       <use style=\"stroke:#000000;stroke-width:0.8;\" x=\"50.14375\" xlink:href=\"#m389b38bd27\" y=\"10.999219\"/>\n      </g>\n     </g>\n     <g id=\"text_12\">\n      <!-- 15.0 -->\n      <g transform=\"translate(20.878125 14.798437)scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-49\"/>\n       <use x=\"63.623047\" xlink:href=\"#DejaVuSans-53\"/>\n       <use x=\"127.246094\" xlink:href=\"#DejaVuSans-46\"/>\n       <use x=\"159.033203\" xlink:href=\"#DejaVuSans-48\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_13\">\n     <!-- perplexity -->\n     <defs>\n      <path d=\"M 41.109375 46.296875 \nQ 39.59375 47.171875 37.8125 47.578125 \nQ 36.03125 48 33.890625 48 \nQ 26.265625 48 22.1875 43.046875 \nQ 18.109375 38.09375 18.109375 28.8125 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 20.953125 51.171875 25.484375 53.578125 \nQ 30.03125 56 36.53125 56 \nQ 37.453125 56 38.578125 55.875 \nQ 39.703125 55.765625 41.0625 55.515625 \nz\n\" id=\"DejaVuSans-114\"/>\n      <path d=\"M 9.421875 75.984375 \nL 18.40625 75.984375 \nL 18.40625 0 \nL 9.421875 0 \nz\n\" id=\"DejaVuSans-108\"/>\n      <path d=\"M 54.890625 54.6875 \nL 35.109375 28.078125 \nL 55.90625 0 \nL 45.3125 0 \nL 29.390625 21.484375 \nL 13.484375 0 \nL 2.875 0 \nL 24.125 28.609375 \nL 4.6875 54.6875 \nL 15.28125 54.6875 \nL 29.78125 35.203125 \nL 44.28125 54.6875 \nz\n\" id=\"DejaVuSans-120\"/>\n      <path d=\"M 9.421875 54.6875 \nL 18.40625 54.6875 \nL 18.40625 0 \nL 9.421875 0 \nz\nM 9.421875 75.984375 \nL 18.40625 75.984375 \nL 18.40625 64.59375 \nL 9.421875 64.59375 \nz\n\" id=\"DejaVuSans-105\"/>\n      <path d=\"M 18.3125 70.21875 \nL 18.3125 54.6875 \nL 36.8125 54.6875 \nL 36.8125 47.703125 \nL 18.3125 47.703125 \nL 18.3125 18.015625 \nQ 18.3125 11.328125 20.140625 9.421875 \nQ 21.96875 7.515625 27.59375 7.515625 \nL 36.8125 7.515625 \nL 36.8125 0 \nL 27.59375 0 \nQ 17.1875 0 13.234375 3.875 \nQ 9.28125 7.765625 9.28125 18.015625 \nL 9.28125 47.703125 \nL 2.6875 47.703125 \nL 2.6875 54.6875 \nL 9.28125 54.6875 \nL 9.28125 70.21875 \nz\n\" id=\"DejaVuSans-116\"/>\n      <path d=\"M 32.171875 -5.078125 \nQ 28.375 -14.84375 24.75 -17.8125 \nQ 21.140625 -20.796875 15.09375 -20.796875 \nL 7.90625 -20.796875 \nL 7.90625 -13.28125 \nL 13.1875 -13.28125 \nQ 16.890625 -13.28125 18.9375 -11.515625 \nQ 21 -9.765625 23.484375 -3.21875 \nL 25.09375 0.875 \nL 2.984375 54.6875 \nL 12.5 54.6875 \nL 29.59375 11.921875 \nL 46.6875 54.6875 \nL 56.203125 54.6875 \nz\n\" id=\"DejaVuSans-121\"/>\n     </defs>\n     <g transform=\"translate(14.798437 103.411383)rotate(-90)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"63.476562\" xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"125\" xlink:href=\"#DejaVuSans-114\"/>\n      <use x=\"166.113281\" xlink:href=\"#DejaVuSans-112\"/>\n      <use x=\"229.589844\" xlink:href=\"#DejaVuSans-108\"/>\n      <use x=\"257.373047\" xlink:href=\"#DejaVuSans-101\"/>\n      <use x=\"317.146484\" xlink:href=\"#DejaVuSans-120\"/>\n      <use x=\"376.326172\" xlink:href=\"#DejaVuSans-105\"/>\n      <use x=\"404.109375\" xlink:href=\"#DejaVuSans-116\"/>\n      <use x=\"443.318359\" xlink:href=\"#DejaVuSans-121\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_23\">\n    <path clip-path=\"url(#p2d36d1aa8e)\" d=\"M 50.14375 16.512093 \nL 54.129464 42.935167 \nL 58.115179 56.987507 \nL 62.100893 65.282264 \nL 66.086607 70.169608 \nL 70.072321 75.033724 \nL 74.058036 82.145916 \nL 78.04375 87.472653 \nL 82.029464 94.272875 \nL 86.015179 100.560163 \nL 90.000893 108.624755 \nL 93.986607 114.047705 \nL 97.972321 119.669452 \nL 101.958036 125.472614 \nL 105.94375 128.403105 \nL 109.929464 132.049486 \nL 113.915179 133.886012 \nL 117.900893 135.802929 \nL 121.886607 137.676763 \nL 125.872321 137.912164 \nL 129.858036 139.014176 \nL 133.84375 138.276408 \nL 137.829464 139.086518 \nL 141.815179 139.455423 \nL 145.800893 139.638589 \nL 149.786607 139.586917 \nL 153.772321 138.918619 \nL 157.758036 139.605128 \nL 161.74375 139.833007 \nL 165.729464 137.844595 \nL 169.715179 139.822567 \nL 173.700893 139.683397 \nL 177.686607 139.898295 \nL 181.672321 139.946739 \nL 185.658036 139.92836 \nL 189.64375 139.896225 \nL 193.629464 139.954303 \nL 197.615179 139.989304 \nL 201.600893 139.96288 \nL 205.586607 139.957506 \nL 209.572321 139.975654 \nL 213.558036 140.004404 \nL 217.54375 139.979525 \nL 221.529464 139.944705 \nL 225.515179 139.997474 \nL 229.500893 139.945315 \nL 233.486607 140.017919 \nL 237.472321 139.930808 \nL 241.458036 140.047439 \nL 245.44375 140.057548 \n\" style=\"fill:none;stroke:#1f77b4;stroke-linecap:square;stroke-width:1.5;\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 50.14375 146.234821 \nL 50.14375 10.334821 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 245.44375 146.234821 \nL 245.44375 10.334821 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 50.14375 146.234821 \nL 245.44375 146.234821 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 50.14375 10.334821 \nL 245.44375 10.334821 \n\" style=\"fill:none;stroke:#000000;stroke-linecap:square;stroke-linejoin:miter;stroke-width:0.8;\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 183.16875 33.012946 \nL 238.44375 33.012946 \nQ 240.44375 33.012946 240.44375 31.012946 \nL 240.44375 17.334821 \nQ 240.44375 15.334821 238.44375 15.334821 \nL 183.16875 15.334821 \nQ 181.16875 15.334821 181.16875 17.334821 \nL 181.16875 31.012946 \nQ 181.16875 33.012946 183.16875 33.012946 \nz\n\" style=\"fill:#ffffff;opacity:0.8;stroke:#cccccc;stroke-linejoin:miter;\"/>\n    </g>\n    <g id=\"line2d_24\">\n     <path d=\"M 185.16875 23.433258 \nL 205.16875 23.433258 \n\" style=\"fill:none;stroke:#1f77b4;stroke-linecap:square;stroke-width:1.5;\"/>\n    </g>\n    <g id=\"line2d_25\"/>\n    <g id=\"text_14\">\n     <!-- train -->\n     <defs>\n      <path d=\"M 34.28125 27.484375 \nQ 23.390625 27.484375 19.1875 25 \nQ 14.984375 22.515625 14.984375 16.5 \nQ 14.984375 11.71875 18.140625 8.90625 \nQ 21.296875 6.109375 26.703125 6.109375 \nQ 34.1875 6.109375 38.703125 11.40625 \nQ 43.21875 16.703125 43.21875 25.484375 \nL 43.21875 27.484375 \nz\nM 52.203125 31.203125 \nL 52.203125 0 \nL 43.21875 0 \nL 43.21875 8.296875 \nQ 40.140625 3.328125 35.546875 0.953125 \nQ 30.953125 -1.421875 24.3125 -1.421875 \nQ 15.921875 -1.421875 10.953125 3.296875 \nQ 6 8.015625 6 15.921875 \nQ 6 25.140625 12.171875 29.828125 \nQ 18.359375 34.515625 30.609375 34.515625 \nL 43.21875 34.515625 \nL 43.21875 35.40625 \nQ 43.21875 41.609375 39.140625 45 \nQ 35.0625 48.390625 27.6875 48.390625 \nQ 23 48.390625 18.546875 47.265625 \nQ 14.109375 46.140625 10.015625 43.890625 \nL 10.015625 52.203125 \nQ 14.9375 54.109375 19.578125 55.046875 \nQ 24.21875 56 28.609375 56 \nQ 40.484375 56 46.34375 49.84375 \nQ 52.203125 43.703125 52.203125 31.203125 \nz\n\" id=\"DejaVuSans-97\"/>\n      <path d=\"M 54.890625 33.015625 \nL 54.890625 0 \nL 45.90625 0 \nL 45.90625 32.71875 \nQ 45.90625 40.484375 42.875 44.328125 \nQ 39.84375 48.1875 33.796875 48.1875 \nQ 26.515625 48.1875 22.3125 43.546875 \nQ 18.109375 38.921875 18.109375 30.90625 \nL 18.109375 0 \nL 9.078125 0 \nL 9.078125 54.6875 \nL 18.109375 54.6875 \nL 18.109375 46.1875 \nQ 21.34375 51.125 25.703125 53.5625 \nQ 30.078125 56 35.796875 56 \nQ 45.21875 56 50.046875 50.171875 \nQ 54.890625 44.34375 54.890625 33.015625 \nz\n\" id=\"DejaVuSans-110\"/>\n     </defs>\n     <g transform=\"translate(213.16875 26.933258)scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-116\"/>\n      <use x=\"39.208984\" xlink:href=\"#DejaVuSans-114\"/>\n      <use x=\"80.322266\" xlink:href=\"#DejaVuSans-97\"/>\n      <use x=\"141.601562\" xlink:href=\"#DejaVuSans-105\"/>\n      <use x=\"169.384766\" xlink:href=\"#DejaVuSans-110\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"p2d36d1aa8e\">\n   <rect height=\"135.9\" width=\"195.3\" x=\"50.14375\" y=\"10.334821\"/>\n  </clipPath>\n </defs>\n</svg>\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "lstm_jax.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.13"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}