nnx_gru.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "bc00d803eb0d4a629024e1847e5e17e4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_2cf1e36c134c4bb5a45e54e6c2ef6a0b",
              "IPY_MODEL_d0d4669753e9437a9f206c436d9d52eb",
              "IPY_MODEL_7b4a616200c94e60b0576d2dc4637d42"
            ],
            "layout": "IPY_MODEL_e63d90dc38f7423d9f5a329771326fa5"
          }
        },
        "2cf1e36c134c4bb5a45e54e6c2ef6a0b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6e32ddf93b48400898bcc3c72072e582",
            "placeholder": "​",
            "style": "IPY_MODEL_2afd5c1a50834ecd98f7ade0107e3ff9",
            "value": "100%"
          }
        },
        "d0d4669753e9437a9f206c436d9d52eb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_db6472197f79495182233114bc8f712c",
            "max": 10000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_dff2a99d8af44a57a26d2c3c9ce0a285",
            "value": 10000
          }
        },
        "7b4a616200c94e60b0576d2dc4637d42": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7ee5b7a3d4cd4d8982f3d64b8dfbeb5d",
            "placeholder": "​",
            "style": "IPY_MODEL_d9c77642bcf5495782aff842711638d5",
            "value": " 10000/10000 [1:23:53&lt;00:00,  2.00it/s, loss=0.0381]"
          }
        },
        "e63d90dc38f7423d9f5a329771326fa5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6e32ddf93b48400898bcc3c72072e582": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2afd5c1a50834ecd98f7ade0107e3ff9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "db6472197f79495182233114bc8f712c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "dff2a99d8af44a57a26d2c3c9ce0a285": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7ee5b7a3d4cd4d8982f3d64b8dfbeb5d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d9c77642bcf5495782aff842711638d5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "%pip install --quiet --upgrade flax transformers\n",
        "%env XLA_PYTHON_CLIENT_MEM_FRACTION=0.975\n",
        "%env TOKENIZERS_PARALLELISM=\"false\"\n",
        "\n",
        "import random\n",
        "import itertools\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import requests\n",
        "import optax\n",
        "import tqdm\n",
        "\n",
        "from typing import Generator\n",
        "from dataclasses import dataclass\n",
        "from pprint import pprint\n",
        "\n",
        "from flax import nnx\n",
        "from transformers import AutoTokenizer, TextStreamer\n",
        "\n",
        "default_dtype = (\n",
        "    jnp.bfloat16\n",
        "    # jnp.float32\n",
        ")\n",
        "\n",
        "@jax.vmap\n",
        "def binary_operator_diag(q_i, q_j):\n",
        "    (A_i, b_i), (A_j, b_j) = q_i, q_j\n",
        "    return A_j * A_i, A_j * b_i + b_j\n",
        "\n",
        "\n",
        "def parallel_scan(a, b, v_0):\n",
        "    a = jnp.concatenate([jnp.zeros((1, *a.shape[1:]), dtype=a.dtype), a])\n",
        "    b = jnp.concatenate([v_0[None], b])\n",
        "    _, v_t = jax.lax.associative_scan(binary_operator_diag, (a, b))\n",
        "    return v_t[1:]\n",
        "\n",
        "\n",
        "class GRUBase(nnx.Module):\n",
        "\n",
        "    def __init__(self, dim: int, rngs: nnx.Rngs):\n",
        "        self.linear_z = nnx.Linear(dim, dim, param_dtype=default_dtype, rngs=rngs)\n",
        "        self.linear_h = nnx.Linear(dim, dim, param_dtype=default_dtype, rngs=rngs)\n",
        "\n",
        "    def __call__(self, x: jax.Array) -> tuple[jax.Array, jax.Array]:\n",
        "        z = nnx.sigmoid(self.linear_z(x))\n",
        "        h_tilde = self.linear_h(x)\n",
        "        return z, h_tilde\n",
        "\n",
        "\n",
        "class ParallelGRU(GRUBase):\n",
        "\n",
        "    @nnx.vmap(in_axes=(None, 0, 0), out_axes=0)\n",
        "    def __call__(self, x: jax.Array, h_prev: jax.Array) -> jax.Array:\n",
        "        z, h_tilde = super().__call__(x)\n",
        "        h = parallel_scan(1 - z, z * h_tilde, h_prev)\n",
        "        return h\n",
        "\n",
        "\n",
        "class SequentialGRU(GRUBase):\n",
        "\n",
        "    def __call__(self, x: jax.Array, h_prev: jax.Array) -> jax.Array:\n",
        "        z, h_tilde = super().__call__(x)\n",
        "        h = (1 - z) * h_prev + z * h_tilde\n",
        "        return h\n",
        "\n",
        "\n",
        "class ResBlock(nnx.Module):\n",
        "\n",
        "    def __init__(self, dim: int, gru: type, *, rngs: nnx.Rngs):\n",
        "        self.gru = gru(dim, rngs=rngs)\n",
        "        self.linear = nnx.Linear(dim, 2 * dim, param_dtype=default_dtype, rngs=rngs)\n",
        "\n",
        "    def __call__(self, x: jax.Array, h_prev: jax.Array) -> tuple[jax.Array, jax.Array]:\n",
        "        h = self.gru(x, h_prev)\n",
        "        x += nnx.glu(self.linear(h))\n",
        "        return x, h\n",
        "\n",
        "\n",
        "class Backbone(nnx.Module):\n",
        "\n",
        "    def __init__(self, depth: int, dim: int, gru: type, *, rngs: nnx.Rngs):\n",
        "\n",
        "        @nnx.split_rngs(splits=depth)\n",
        "        @nnx.vmap\n",
        "        def create_block(rngs: nnx.Rngs):\n",
        "            return ResBlock(dim, gru, rngs=rngs)\n",
        "\n",
        "        self.layers = create_block(rngs)\n",
        "        self.depth = depth\n",
        "        self.dim = dim\n",
        "\n",
        "    def __call__(self, x: jax.Array, h_prev: jax.Array) -> tuple[jax.Array, jax.Array]:\n",
        "\n",
        "        @nnx.scan(in_axes=(nnx.Carry, 0, 0), out_axes=(nnx.Carry, 0))\n",
        "        def forward(x, h_prev, backbone):\n",
        "            return backbone(x, h_prev)\n",
        "\n",
        "        return forward(x, h_prev, self.layers)\n",
        "\n",
        "\n",
        "class Model(nnx.Module):\n",
        "\n",
        "    def __init__(self, vocab_size: int, depth: int, dim: int, gru: type, *, rngs: nnx.Rngs):\n",
        "        self.embeddings = nnx.Embed(vocab_size, dim, param_dtype=default_dtype, rngs=rngs)\n",
        "        self.backbone = Backbone(depth, dim, gru, rngs=rngs)\n",
        "        self.linear_out = nnx.Linear(dim, vocab_size, param_dtype=default_dtype, rngs=rngs)\n",
        "        self.vocab_size = vocab_size\n",
        "        self.depth = depth\n",
        "        self.dim = dim\n",
        "\n",
        "    def __call__(self, x: jax.Array, h_prev: jax.Array) -> tuple[jax.Array, jax.Array]:\n",
        "        x = self.embeddings(x)\n",
        "        x, h = self.backbone(x, h_prev)\n",
        "        x = self.linear_out(x)\n",
        "        return x, h\n",
        "\n",
        "\n",
        "class ParallelModel(Model):\n",
        "    def __init__(self, *, vocab_size: int, depth: int, dim: int, rngs: nnx.Rngs):\n",
        "        super().__init__(vocab_size, depth, dim, ParallelGRU, rngs=rngs)\n",
        "\n",
        "\n",
        "class SequentialModel(Model):\n",
        "    def __init__(self, *, vocab_size: int, depth: int, dim: int, rngs: nnx.Rngs):\n",
        "        super().__init__(vocab_size, depth, dim, SequentialGRU, rngs=rngs)\n",
        "\n",
        "\n",
        "def to_sequential(parallel_model: ParallelModel) -> SequentialModel:\n",
        "    sequential_model = nnx.eval_shape(\n",
        "        lambda: SequentialModel(\n",
        "            vocab_size=parallel_model.vocab_size,\n",
        "            depth=parallel_model.depth,\n",
        "            dim=parallel_model.dim,\n",
        "            rngs=nnx.Rngs(0)\n",
        "        )\n",
        "    )\n",
        "    _, state = nnx.split(parallel_model)\n",
        "    nnx.update(sequential_model, state)\n",
        "    return sequential_model\n",
        "\n",
        "\n",
        "@nnx.jit\n",
        "def train_step(model: ParallelModel, optimizer: nnx.Optimizer, batch: jax.Array) -> jax.Array:\n",
        "    inputs, targets = batch[:, :-1], batch[:, 1:]\n",
        "    initial_state = jnp.zeros((model.depth, inputs.shape[0], model.dim), default_dtype)\n",
        "\n",
        "    @nnx.value_and_grad\n",
        "    def loss_fn(model):\n",
        "        logits, _ = model(inputs, initial_state)\n",
        "        loss = optax.softmax_cross_entropy_with_integer_labels(logits, targets).mean()\n",
        "        return loss\n",
        "\n",
        "    loss, grads = loss_fn(model)\n",
        "    optimizer.update(grads)\n",
        "    return loss\n",
        "\n",
        "\n",
        "@nnx.jit\n",
        "def prefill(model: ParallelModel, tokens: jax.Array) -> tuple[jax.Array, jax.Array]:\n",
        "    initial_state = jnp.zeros((model.depth, tokens.shape[0], model.dim), default_dtype)\n",
        "    logits, state = model(tokens, initial_state)\n",
        "    final_logits = logits[..., -1, :]\n",
        "    final_state = state[..., -1, :]\n",
        "    return final_logits, final_state\n",
        "\n",
        "\n",
        "@nnx.jit\n",
        "def sample(\n",
        "    model: SequentialModel,\n",
        "    token: jax.Array,\n",
        "    state: jax.Array,\n",
        "    *,\n",
        "    temperature: float,\n",
        "    rng: jax.Array,\n",
        ") -> tuple[jax.Array, jax.Array]:\n",
        "    logits, state = model(token, state)\n",
        "    token = jax.random.categorical(rng, logits / temperature)\n",
        "    return token, state"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uYQurJm1ZSww",
        "outputId": "345a9b53-b771-4978-ead9-80f209e70445"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: XLA_PYTHON_CLIENT_MEM_FRACTION=0.975\n",
            "env: TOKENIZERS_PARALLELISM=\"false\"\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "texts = []\n",
        "for url in (\n",
        "    \"https://www.gutenberg.org/ebooks/98.txt.utf-8\",      # a tale of two cities\n",
        "    \"https://www.gutenberg.org/ebooks/2701.txt.utf-8\",    # moby dick\n",
        "    \"https://www.gutenberg.org/ebooks/2554.txt.utf-8\",    # crime and punishment\n",
        "    \"https://www.gutenberg.org/ebooks/4300.txt.utf-8\",    # ulysses\n",
        "    \"https://www.gutenberg.org/ebooks/768.txt.utf-8\",     # wuthering heights\n",
        "    \"https://www.gutenberg.org/ebooks/132.txt.utf-8\",     # the art of war\n",
        "    \"https://www.gutenberg.org/ebooks/76.txt.utf-8\",      # huckleberry finn\n",
        "    \"https://www.gutenberg.org/cache/epub/100/pg100.txt\", # complete works of shakespeare\n",
        "):\n",
        "    response = requests.get(url)\n",
        "    texts.append(response.text)"
      ],
      "metadata": {
        "id": "b80cTd7vh3Df"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\", clean_up_tokenization_spaces=True)\n",
        "tokens = tokenizer(texts, return_tensors=\"np\").input_ids\n",
        "print(f\"{sum(t.size for t in tokens)/1_000_000:.1f}M tokens\")\n",
        "\n",
        "def gen_sample(length: int) -> Generator[np.ndarray, None, None]:\n",
        "    while True:\n",
        "        book = random.choice(tokens)\n",
        "        start = random.choice(range(book.size - length))\n",
        "        stop = start + length\n",
        "        sample = book[start:stop]\n",
        "        yield sample\n",
        "\n",
        "def gen_batch(batch_size: int, length: int) -> Generator[jax.Array, None, None]:\n",
        "    sampler = iter(gen_sample(length))\n",
        "    while True:\n",
        "        batch = jnp.stack(list(itertools.islice(sampler, batch_size)))\n",
        "        yield batch"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "onrQ-MvViErr",
        "outputId": "7fef197f-10af-41cc-e8cd-1bed90ec35e2"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n",
            "Token indices sequence length is longer than the specified maximum sequence length for this model (228351 > 1024). Running this sequence through the model will result in indexing errors\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3.8M tokens\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "@dataclass\n",
        "class Config:\n",
        "    seed: int = 47\n",
        "    depth: int = 16\n",
        "    dim: int = 1536\n",
        "    vocab_size: int = tokenizer.vocab_size\n",
        "    init_learning_rate: float = 1e-4\n",
        "    peak_learning_rate: float = 1e-3\n",
        "    end_learning_rate: float = 1e-5\n",
        "    warmup_steps: int = 100\n",
        "    decay_steps: int = 10_000\n",
        "    weight_decay: float = 0\n",
        "    gradient_clipping: float = 1.0\n",
        "    batch_size: int = 80\n",
        "    length: int = 512\n",
        "\n",
        "\n",
        "def initialize(cfg: Config) -> tuple[ParallelModel, nnx.Optimizer]:\n",
        "    parallel_model = ParallelModel(\n",
        "        vocab_size=cfg.vocab_size,\n",
        "        depth=cfg.depth,\n",
        "        dim=cfg.dim,\n",
        "        rngs=nnx.Rngs(cfg.seed),\n",
        "    )\n",
        "    tx = optax.chain(\n",
        "        optax.clip_by_global_norm(cfg.gradient_clipping),\n",
        "        optax.adamw(\n",
        "            learning_rate=optax.warmup_cosine_decay_schedule(\n",
        "                init_value=cfg.init_learning_rate,\n",
        "                peak_value=cfg.peak_learning_rate,\n",
        "                end_value=cfg.end_learning_rate,\n",
        "                warmup_steps=cfg.warmup_steps,\n",
        "                decay_steps=cfg.decay_steps,\n",
        "            ),\n",
        "            weight_decay=cfg.weight_decay,\n",
        "        )\n",
        "    )\n",
        "    optimizer = nnx.Optimizer(parallel_model, tx)\n",
        "    params = nnx.split(parallel_model, nnx.Param)[1]\n",
        "    print(f\"{sum(jax.tree.flatten(jax.tree.map(jnp.size, params))[0])/1_000_000:.1f}M params\")\n",
        "    return parallel_model, optimizer\n",
        "\n",
        "\n",
        "cfg = Config()\n",
        "pprint(cfg)\n",
        "parallel_model, optimizer = initialize(cfg)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GIvxJe2Thwkx",
        "outputId": "7951de90-148a-4610-8aed-1388b406b45b"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Config(seed=47,\n",
            "       depth=16,\n",
            "       dim=1536,\n",
            "       vocab_size=50257,\n",
            "       init_learning_rate=0.0001,\n",
            "       peak_learning_rate=0.001,\n",
            "       end_learning_rate=1e-05,\n",
            "       warmup_steps=100,\n",
            "       decay_steps=10000,\n",
            "       weight_decay=0,\n",
            "       gradient_clipping=1.0,\n",
            "       batch_size=80,\n",
            "       length=512)\n",
            "305.5M params\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dataloader = gen_batch(cfg.batch_size, cfg.length)\n",
        "progbar = tqdm.notebook.trange(cfg.decay_steps, mininterval=1)\n",
        "losses = []\n",
        "try:\n",
        "    for step, batch in zip(progbar, dataloader):\n",
        "        loss = train_step(parallel_model, optimizer, batch)\n",
        "        progbar.set_postfix(loss=loss.item(), refresh=False)\n",
        "        losses.append(loss.item())\n",
        "finally:\n",
        "    plt.semilogy(losses)\n",
        "    plt.grid(True, which=\"both\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 462,
          "referenced_widgets": [
            "bc00d803eb0d4a629024e1847e5e17e4",
            "2cf1e36c134c4bb5a45e54e6c2ef6a0b",
            "d0d4669753e9437a9f206c436d9d52eb",
            "7b4a616200c94e60b0576d2dc4637d42",
            "e63d90dc38f7423d9f5a329771326fa5",
            "6e32ddf93b48400898bcc3c72072e582",
            "2afd5c1a50834ecd98f7ade0107e3ff9",
            "db6472197f79495182233114bc8f712c",
            "dff2a99d8af44a57a26d2c3c9ce0a285",
            "7ee5b7a3d4cd4d8982f3d64b8dfbeb5d",
            "d9c77642bcf5495782aff842711638d5"
          ]
        },
        "id": "xfTr-NxoiVtc",
        "outputId": "ffe65bbd-158e-429f-9848-1182fa303bbc"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/10000 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "bc00d803eb0d4a629024e1847e5e17e4"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAi4AAAGdCAYAAAA1/PiZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABT/UlEQVR4nO3deVxU9f4/8NfMMGzKjoAoiAtuuG+ImmluqTfLtluaedsXK83K22bW9WfacvuWReXtVtYtlxaXFjMR911R3BVNcQEBUWHYhFnO7w9kYJgBZmDmnDlnXs/Ho0fMOZ858543A7z9nM+iEgRBABEREZEMqKUOgIiIiMheLFyIiIhINli4EBERkWywcCEiIiLZYOFCREREssHChYiIiGSDhQsRERHJBgsXIiIikg0vqQNwNpPJhOzsbAQEBEClUkkdDhEREdlBEAQUFRUhOjoaanXd/SqKK1yys7MRExMjdRhERETUCBcuXEDr1q3rPK+4wiUgIABA5RsPDAx02nX1ej3WrVuH0aNHQ6vVOu26ZI25FgfzLA7mWRzMszhcmWedToeYmBjz3/G6KK5wqbo9FBgY6PTCxd/fH4GBgfyhcDHmWhzMsziYZ3Ewz+IQI88NDfPg4FwiIiKSDRYuREREJBssXIiIiEg2WLgQERGRbLBwISIiItlg4UJERESywcKFiIiIZIOFCxEREckGCxciIiKSDRYuREREJBssXIiIiEg2WLgQERGRbLhl4TJx4kSEhITg7rvvljoUszWHc/DzWTW2nMqXOhQiIiKP5ZaFy/Tp0/Htt99KHYaFnWevYkuOGgcvFEodChERkcdyy8Jl2LBhCAgIkDoMC96aylRVGE0SR0JEROS5nF64bNmyBbfddhuio6OhUqmwatUqqzbJycmIi4uDr68vEhMTsWfPHmeH4XTeXjcKFwMLFyIiIqk4vXApKSlBz549kZycbPP88uXLMXPmTMyZMwf79+9Hz549MWbMGOTl5Tk7FKfy8WKPCxERkdS8nH3BsWPHYuzYsXWe/+CDD/DYY4/hoYceAgB8/vnn+P333/HVV1/h5Zdfdvj1ysvLUV5ebn6s0+kAAHq9Hnq93uHr1UWjEgAA1ysMTr0uWavKL/PsWsyzOJhncTDP4nBlnu29ptMLl/pUVFQgLS0Nr7zyivmYWq3GyJEjsXPnzkZdc/78+Xjrrbesjq9btw7+/v6NjrW2c9kqABpkXsjCmjUXnHZdqltKSorUIXgE5lkczLM4mGdxuCLPpaWldrUTtXDJz8+H0WhEZGSkxfHIyEicOHHC/HjkyJE4ePAgSkpK0Lp1a/z4449ISkqyec1XXnkFM2fOND/W6XSIiYnB6NGjERgY6LTYc7adxepzp9AiMgrjxvVy2nXJml6vR0pKCkaNGgWtVit1OIrFPIuDeRYH8ywOV+a56o5JQ0QtXOy1fv16u9v6+PjAx8fH6rhWq3VqUn29K1NlFMAfCpE4+3tItjHP4mCexcE8i8MVebb3eqJOhw4PD4dGo0Fubq7F8dzcXERFRYkZisO0N6ZD6zk4l4iISDKiFi7e3t7o27cvUlNTzcdMJhNSU1PrvBXkLrQaFQBAbxQkjoSIiMhzOf1WUXFxMU6fPm1+fPbsWaSnpyM0NBSxsbGYOXMmpk6din79+mHAgAH48MMPUVJSYp5l5CzOnlWkRmVPS4XeyFHrLsbZAeJgnsXBPIuDeRaHO8wqUgmC4NQuhE2bNmH48OFWx6dOnYrFixcDAD755BO89957yMnJQa9evbBw4UIkJiY26XWTk5ORnJwMo9GIjIwMLFmyxKmzig5eUeGrDA3aBgiY0c3otOsSERFR5ayiSZMmobCwsN7JNU4vXKSm0+kQFBSE/Px8p84qWn/0Ep5adhjdowOw4in3vq0ld5wdIA7mWRzMsziYZ3G4elZReHh4g4WLW84qcgZnj3j28am8lt7EWUVi4ewAcTDP4mCexcE8i8NjZhXJmTdnFREREUmOhYudqqZDGziriIiISDIsXOxUPR2aPS5ERERSUewYF2dPh4bpxnRoo4nT7VyM0xrFwTyLg3kWB/MsDkVOh5aKq6dD55QC8w96wd9LwPz+nA5NRETkTJwO7eTp0H/l6nDrJ7vg763BwdkjnHZdssZpjeJgnsXBPIuDeRYHp0O7kLOnavn7egMADCaBPxQi4bRGcTDP4mCexcE8i4PToWXAS109OFdhnVRERESywcLFTlXToQUBMJpYuBAREUmBhYudqqZDA9whmoiISCqKHePi9OnQQvVMorLycnipeA/VVTitURzMsziYZ3Ewz+LgdGgncvV0aEEAZuyqrPPm9TOgOesWIiIip+F0aCdPh9br9eg2dwNMggpbXxqKqEBfp12bLHFaoziYZ3Ewz+JgnsXB6dAu5IqpWl4qoEIAoNLwB0MEnNYoDuZZHMyzOJhncXA6tExUmCoH6H6QkiFxJERERJ6JhUsjrDyQhaLrHABGREQkNhYujTTv9+NSh0BERORxWLg00rK9F6QOgYiIyOOwcHFAK39FTcAiIiKSHcXOKnL2AnR6vR6zehoxfWd1yg6ev4KuLZ035ZoqcSEpcTDP4mCexcE8i4ML0DmRqxegq1KzcAn2FvBWX2M9rYmIiMgeXIDOBQvQpaSk4E9dNNYczTMfPzV3tNNegypxISlxMM/iYJ7FwTyLgwvQuZCrFiF6ZngHi8KFPyCuw4WkxME8i4N5FgfzLA4uQCcj8ZHNLR4fv6STKBIiIiLPw8KliZ78Lk3qEIiIiDwGC5dGmJwYa/76/NVSCSMhIiLyLCxcGmFKUhvz14IA5OquSxgNERGR52Dh0gidoyxHO3+wjpsuEhERiYGFixNcK62QOgQiIiKPwMLFCY5xZhEREZEoFLuOiyuW/K/5/wOv34Le/28DAODitTIuM+1EXLpbHMyzOJhncTDP4uCS/04k1pL/NdVc/v+jJINLX4uIiEjJuOS/i5b8r7nMcfzsdebzR+aMhI8X77w5A5fuFgfzLA7mWRzMszi45L8LuWrZ55rXvbNPK6zYnwUA2HHmGkYnRDn99TwZl+4WB/MsDuZZHMyzOLjkv0w9mBRn/vrHtIvSBUJEROQhWLg0QUyIn/nrC1xBl4iIyOVYuDRBWHMf89cncopYvBAREbkYCxcnulTIpf+JiIhciYVLE3VrVT3yueg61w8gIiJyJRYuTfTFg/3MXz/yzT4JIyEiIlI+Fi5N1DLIz+Ixx7kQERG5DgsXJ7vp3Y1Sh0BERKRYil2AztV7FdXUKyYI6RcKzY/3nrmMXjHBTnttT8M9R8TBPIuDeRYH8ywO7lXkRFLsVVTldCHw8bHqGvDhjkb0DFNEWomIiETBvYpE2KuoiiAI6PhGisWxU3NHO+21PQ33HBEH8ywO5lkczLM4uFeRC4mxV1FD7ahpuOeIOJhncTDP4mCexcG9ioiIiIjswMLFSb5+qL/F478v2ilRJERERMrFwsVJhneKgL+3xvx499mruFJcLmFEREREysPCxYk2vDDM4nGZ3ihNIERERArFwsWJvDQqi8dD3uFidERERM7EwsWJfLWahhsRERFRo7FwcaLmPtazyzPzSySIhIiISJlYuDhZiL/lPPRh72+SJhAiIiIFYuHiZPPv7CF1CERERIrFwsXJxiREWh2bv+a4BJEQEREpDwsXJ1OpVFbHFm05g4vXSiWIhoiISFlYuLjA8X/danXs001/SRAJERGRsrBwcQFvL+u0lutNEkRCRESkLIrdHVqv10Ov1zv1ejX/Xx9BEKyO/bz/Ioa0D8HferR0WkxK5UiuqfGYZ3Ewz+JgnsXhyjzbe02VYOuvrAwlJycjOTkZRqMRGRkZWLJkCfz9/SWLZ/pO2zXhR0kGkSMhIiJyf6WlpZg0aRIKCwsRGBhYZzvFFC5VdDodgoKCkJ+fX+8bd5Rer0dKSgpGjRoFrVbbYPvZvxzDsr0XrY6fmjvaaTEplaO5psZhnsXBPIuDeRaHK/Os0+kQHh7eYOGi2FtFWq3WJR9ee69rqqMcLLhuQosAHydHpUyu+h6SJeZZHMyzOJhncbgiz/Zej4NzXcRgtF259J+3XuRIiIiIlIOFi4tEB/tJHQIREZHisHBxkaeGtcd9/WNsnot7+XfojZweTURE5CgWLi7SzMcLC+6qe9+iZXvOixgNERGRMrBwkUh24XWpQyAiIpIdFi4SUdYkdCIiInGwcHExf2+NzeMCWLkQERE5ioWLi+1+dYTN47vOXBU5EiIiIvlj4eJiAb62F9Q5eKFA3ECIiIgUgIULERERyQYLFxG0qmMxujOXi0WOhIiISN5YuIjghyeTkNQuzOr4tdIKCaIhIiKSLxYuImgV7Ieljw+0Ov7KisPIyC2SICIiIiJ5YuEioYzcYoz+vy1Sh0FERCQbLFzcgNHENV2IiIjswcLFDSzfe0HqEIiIiGSBhYsbeHXlYe4WTUREZAcWLiJ69+66d4tesf+iiJEQERHJEwsXEd3bL6bONV0Ky/QiR0NERCQ/XlIH4Cp6vR56vfOKgaprNfWaKTMGo+ub662On8svcWq8cuasXFP9mGdxMM/iYJ7F4co823tNlSAIipjSkpycjOTkZBiNRmRkZGDJkiXw9/eXOiybpu+0XS9+lGQQORIiIiL3UFpaikmTJqGwsBCBgYF1tlNM4VJFp9MhKCgI+fn59b5xR+n1eqSkpGDUqFHQam1vnGiv+NnrbB7fPutmRAT4NOnaSuDMXFPdmGdxMM/iYJ7F4co863Q6hIeHN1i4KPZWkVardcmH11XXBYCh72/B6XljoVKpXHJ9uXFlrqka8ywO5lkczLM4XJFne6/HwbluxGgS0OtfKVKHQURE5LZYuLgZzi4iIiKqGwsXCcy6tVO955/43z78dihbpGiIiIjkg4WLBAa2C6v3/J9Hc/HMkgMiRUNERCQfLFwkoKx5XEREROJh4SKBduHN7Gr360HeLiIiIqqJhYsEQpp5Y9OLwxps9+xS3i4iIiKqiYWLROLCm+Gj+3pJHQYREZGssHCR0O29WkkdAhERkaywcJGYveNdiIiIiIWL5GJC3XMjSCIiInfEwsXNHTh/DWM/2oodf+VLHQoREZHkWLhIrKElXSZ9sRvHL+kw6YvdosRDRETkzli4SKxzVEC958v0RpEiISIicn8sXCQ2Y2S81CEQERHJBgsXifl7eyG0mbfUYRAREckCCxciIiKSDRYuMnKd412IiMjDsXCRkc6z1yKn8LrUYRAREUmGhYvM/Lz/otQhEBERSYaFCxEREckGCxeZqTCYpA6BiIhIMixc3MADA9vY3bawTO/CSIiIiNwbCxc38NwtHbDs8YGIDvJtsO3iHZm4eK1UhKiIiIjcDwsXN+ClUWNguzD4emvsaj/knY0wmRra5YiIiEh5WLi4EwdqkTs+3e66OIiIiNwUCxc34kgfyqGLhThzudhlsRAREbkjFi5uRBAcu/1zy783Y+me8y6KhoiIyP2wcHEjfduEOvycTzeddkEkRERE7omFixuZM6ErHhoc16jnHsvWYXV6lnMDIiIicjNeUgdA1QJ9tXh9fFd8vT3T7udcuFoGABi3cCsAIKyZD4bEh7siPCIiIsm5ZY/Lb7/9hk6dOiE+Ph7//e9/pQ5HVBq1Cmueu8mh5wx9d6P56xM5OmeHRERE5DbcrsfFYDBg5syZ2LhxI4KCgtC3b19MnDgRYWFhUocmGm8vx+rJ81e5IB0REXkGt+tx2bNnDxISEtCqVSs0b94cY8eOxbp166QOS1ShzbylDoGIiMgtOb1w2bJlC2677TZER0dDpVJh1apVVm2Sk5MRFxcHX19fJCYmYs+ePeZz2dnZaNWqlflxq1atkJXlWYNOWbgQERHZ5vRbRSUlJejZsycefvhh3HnnnVbnly9fjpkzZ+Lzzz9HYmIiPvzwQ4wZMwYnT55ERESEw69XXl6O8vJy82OdrnKMh16vh17vvA0Jq67lzGvWJ+3V4ej79saGG9ZiMplEi9FVxM61p2KexcE8i4N5Focr82zvNVWCo6ueOUClUmHlypW44447zMcSExPRv39/fPLJJwAq/9DGxMTg2Wefxcsvv4wdO3bgvffew8qVKwEAM2bMwIABAzBp0iSbr/Hmm2/irbfesjq+ZMkS+Pv7O/9NiWj6TsfryolxRgxryX2MiIhIXkpLSzFp0iQUFhYiMDCwznaiFi4VFRXw9/fHTz/9ZFHMTJ06FQUFBVi9ejUMBgO6dOmCTZs2mQfn7tixo87BubZ6XGJiYpCfn1/vG3eUXq9HSkoKRo0aBa1W67Tr1id+tuNje14d2wkPDWrjgmjEI0WuPRHzLA7mWRzMszhcmWedTofw8PAGCxdRZxXl5+fDaDQiMjLS4nhkZCROnDhRGZCXF/79739j+PDhMJlMmDVrVr0zinx8fODj42N1XKvVuuTD66rrOotGo3Hr+Bzh7rlWCuZZHMyzOJhncbgiz/Zez+2mQwPAhAkTMGHCBKnDkFxEgA/yisobbljD3N+O4dDFAiy4swf8vDXm4ydzivDz/ot46ub2COHgXyIikilRC5fw8HBoNBrk5uZaHM/NzUVUVJSYochC6xA/hwsXAFidno0LV0ux4unB5mNjPtwCAMgqKEPypD5Oi5GIiEhMohYu3t7e6Nu3L1JTU81jXEwmE1JTU/HMM8849bXkPqsIABZMTMCMHw7h2KUih5+7/3wB9Ho9dGV6eGlU5uNHLha6/ah7zg4QB/MsDuZZHMyzOBQ5q6i4uBinT1fuWNy7d2988MEHGD58OEJDQxEbG4vly5dj6tSpWLRoEQYMGIAPP/wQP/zwA06cOGE19sURycnJSE5OhtFoREZGhiJmFVVpzOwiAIhrLiCzWGVxLNxXwOzeRmeERURE5DSSzSratGkThg8fbnV86tSpWLx4MQDgk08+wXvvvYecnBz06tULCxcuRGJiolNeX6fTISgoSBGziqo0ZnZRXdqE+mP980Ocdj1X4OwAcTDP4mCexcE8i0ORs4qGDRuGhmqhZ555xum3hmpT0qyiWbd2wrtrTzrlWiqV/SO3pcbZAeJgnsXBPIuDeRaHlLOK3G6vIrI2qkvjb6HVlnmlFCaTgKe+S8P7fzqnGCIiIhILCxcZiI8McOr19mRexR9HcvDJxtNOvS4REZGrsXCRicwF4512ret6Ds4lIiJ5cssF6JxBCdOha+vaMqBRU6NrMxgM5q9rvp/MKyXw1qgRHezX5NdoCnfItSdgnsXBPIuDeRaHIqdDS0XJ06GrXDcC/9zT9Frzic5GLDpRuaruR0mVRUyZAXh5r5fFMSIiIrG4xSaLUlDidOgqpRUG9Jy7ocnX+WJKbzz2vwMAgFNzRwMATuQU4bbknRbHpOIOufYEzLM4mGdxMM/iUOR0aHehpOnQVXycNCTJS1P9ba96Lxobx6TGaY3iYJ7FwTyLg3kWB6dDk6j+b32GxePzV0phNFV3vK09cknskIiIiOyi2B4XJVKpGm5jj0MXC81fD39/E87ml6BlkK/52JPf7ce+10civLmPxfOulVQgq6AM3VoFOScQIiIiB7HHRUbUzqpcajibXwIAuFR43eJ4Qan16O7E+an428fbsP/8NZvXMpkEGIwmp8dIRERUhYWLjDi/bKmbrTHbFYbKomRLxmWbz7nj0+0Y8s5GczsiIiJnU+ytIiWu41JzHAoANPPWoKTCNYvJ5ReVISbYBxq1dblkMpps5qHqFtSxrGtIiG78jC53yLUnYJ7FwTyLg3kWB9dxcSJPWMdFEIAZuyprzYlxRmzLUePyddf1w8Q1F/B89+rCaPrOyte+tbURY2OsPzZV51/qYUDrZi4Li4iIFIjruChwHRcAiJ+9DgCw9aWhmPp1Gs7cGKPiKjXXdKl67eeGt8ezt7SvM7bVTw9E15ZN63Fxh1wrHfMsDuZZHMyzOLiOiwspcR0XAPjqH/1QXG5ETFgA1DZu4zibVqtFucGIn9Iumo+pNep6c+Dl5eWUHEmda0/BPIuDeRYH8ywOKddxUWzholS3dI40fx3s5/ofzjzddSzbewEfpFSv/bI386rLX5eIiMgWFi4yFtrM2+WvMeDtVKtj209fMX+tN5qg1XByGhERiYN/cajRtp66jPjX/sA3OzItpk+rRJ24TUREnoQ9LtQoJpOAKV/uAQDM+eUoFm3+S+KIiIjIE7DHhRpl3prjFo+za6y864IFfomIiAAouMdFiQvQ1eanla7u/HLb2TrPGQyGJuXJHXOtRMyzOJhncTDP4uACdE7kCQvQ1VZYAfz3hAa9wkz45bxG6nDMZvUwoFUDC9DllgFhPoAX+/yIiAhcgE6xC9DZcrWkAokLNkkdhtmv05LQLrwZvNQqm2vNpB7Pw5NL0tGjVSB+fnKg1Xl3zrWSMM/iYJ7FwTyLgwvQuZBSF6CzxcvLvTY1rDCp0H/+RnSKCsDKpwdbnf/pQDYA4FCWrt5cumOulYh5FgfzLA7mWRxSLkDHjnoFcLcus91nr6C0wogD5wtsnq/Zx5dVUIbCMj2+23UOV0sqmvzaq9OzcCSrsMnXISIi96TYHhdP4m43+8oa2LE69USe+es5q48CELD+eB7eXnMch98c0+jX3X3mCqYvSwcAZC4Y3+jrEBGR+2KPiwIE+LpX/fnxhtNWx45l63D3Zzuw56zldgGlFQasP55342sj/vXrUYvzjgzBysgrbkS0REQkJyxcFMBXq8Ef02+SOox6jVu4FfvOXcO9i3bW2+6bnecAACV6YG/mNfT9f+uxOj1LjBCJiEgG3Ouf6tRoXVo6bwaVM607moPU43kNN6zhcFYhXt3nBezbCwCYviwdt/dqZdHm/JVS+HqrERHgC6Byz6TZq444J2giInJb7HFRkLm3J0gdgpXH/5eG5fsu1Hl+x19XrI59s/N8vde8VlKBoe9txIB51RtA/px2sfFBEhGRbCi2x8UTVs6t7b5+rTB79dGGG7o7G+Naaub9VE6h1fHcwrI625NtcvhMKwHzLA7mWRxcOdeJPHHlXFum75R/LTqghQl7Llt2Bs7vb4D/jbd2tgj48Ejlg4+SDACAdRdV+P1C9erBVceJiEgeuHKuB62cW1P87HVSh9Bkd/SMwqqDOVbHf38mCR0jA3DgfAHu/aJyZ+pTc0cDABZuOI2PN54xt606TnWTy2da7phncTDP4uDKuS7kSSvnKs3GjHybx3/cfwlvTgiFxqv6Y6vValFaYbAoWqqOk334mRYH8ywO5lkcXDmXnO6RIW1xet5YqcNolMIy27d5THV0Do77aKsrwyEiIjfCwkVhWgZVTg++s0+rBlrKT4XBBKPJsng5lq1D5pVSiSIiIiKxKfZWkadKfeFm5BdVIDbM3+qPvNwt23sBu89exfv39DQfG7eQvS1ERJ6EhYvC+Ht7ITZMud/Ws/klWHfUeuAuERF5Bt4qUjCV1AG4yKItZxps89dl7ltERKRELFxIkZ74X5rUIRARkQuwcFEwlVK7XOyQq7veqOfpjSZcLip3cjREROQsLFyIarjt423oP289TuUWNdj2+CUd3vr1KK6WVIgQGRERASxcFE2lUuHTyX3w7t09pA5FfELlZoxTvtyN1elZdj/tRE5lwfL74UsNth370VZ8vT0Tr6443OgwiYjIMYqdfuKJmyzaMqpzOABg1k+HJI5EXAKAf687ga2n8rH1VD7GJUQ49Hyj0Wj39/podqHsPheAfD/TcsM8i4N5Fgc3WXQibrJYPyVsvugIX42ArsEC9l+p7FSsa9NFQbAcC1SVpzGtTRgXY6r3NarahvkIeKOP0QlRExF5Lm6y6KGbLNZFCZsvOqpNqD/OXa1cVTf5/p7o2ToIkYG+5vMfpp7GT/uzsPLJgWgR4AOgOk/ThrXDjBEd6r1+VdvYUD+kPn+TK96CS8n9My0XzLM4mGdxcJNFF+Imi1RVtADAtKUHAQCZC8abjyVvqlwP5j/bzuHNCQkWz9Wo1fZv+KVSyfozwc+0OJhncTDP4uAmi0QiueX9TUg7d83imK1OR0V1QxIRKQgLF/IoZ/JLMOXL3RbHVCoVVqdn4fVV1bODlHUDlYhIORR7q4jq1rdNiFWvgycprbAcSFtcbsD0ZemNvp7Kk1f6IyISGXtcPFD3VkFSh+BWfkq7aHVM4M0iIiK3xMKFPNKO0/lOuxb7W4iIxMPChTzSNzsz6z1f1xiXwxcLkbzxNCoM9a/xQkRErsExLuSR/jyaW+/5um4U3fbJNgCVU6CJiEh87HEhsqGqx0UQBHy57Sz2ZV61OH8iR1f9gDUMEZFoWLh4oF4xwXhhVEcAwPv39JQ4Gvf01fazuHitFH8ezcHc347h7s93Sh0SERGBt4o8xoqnB+GJ/6Xhbz1aYkLPaKjVKjw4KA5Bflq8+ONBc7tAXy/ortve18eTVBhMGPLORrvassOFiEg87HHxEH1iQ7D3tZGYc1sC1OrKP7VBfpXLKz86pC0A4KUxnTCxdyvJYpQTFitERNJg4UJ4bXwXbH5pGJ4e1l7qUIiIiOql2FtFer0eer3eqder+X+liQ70hsFggMnEab51GfZe9a2ji9dKLc7Z+lyUG0woKTcgtJm3y2NrDKV/pt0F8ywO5lkcrsyzvddUCbZ2mJOh5ORkJCcnw2g0IiMjA0uWLIG/v7/UYcnOz2fV2JLDjjhHRPoJeLWX0er4nDQNCipUeKuPAcE+EgRGRCQjpaWlmDRpEgoLCxEYGFhnO8UULlV0Oh2CgoKQn59f7xt3lF6vR0pKCkaNGqXoLdPn/n4C3+46L3UYstKhRTP88dxgq+Pxs9cBAOZPTMDdfdxv7JCnfKalxjyLg3kWhyvzrNPpEB4e3mDhothbRVqt1iUfXldd110M7RjBwsVBFUYBRqix6WQewpv7oF9cqMV5jVpj/szM+ukgyg0mfPj3Xm6zOaPSP9PugnkWB/MsDlfk2d7rKbZwocYZ0SUC3z48AJ2iApD4dqrU4cjC+aul6Dx7rflx5oLxNtuVVhjww77KDR3/eWtnRAf7iRIfEZGScDADWVCpVBjasQUiA33hp9VIHY4sDXw7FTmF162Om4SaX1c+qDCYoLC7tURELsXCheq08cVhUocgSzm663h7zXGr47ULlNIKA/rOTcE9XJWXiMhuLFyoTlFBvnhxdEepw5ClCzWmS8/6+RCuFJdbtdl15gqKyg3Yd+6amKEREckaCxeql7939TCoN/7Wtd62z42Id3U4snHgfIHF477/bz3eWH3U/Pi7XedhrGfJnHNXSpC88TSKrnNNCiKimjg4l+pV8+ZGQnT908vbhTdzbTAyt/JAlvnrzzf/VW/bMR9uwXW9CWfzS7gRJhFRDexxoXrVHJfR0PRdEweZOs11fWV3zE9pF9nrQkRUAwsXqtfwzhEAgPDmDS/96ibLkijOA1/ukToEIiK3wVtFVK/2LZpj+8u3INTfG4ezCutt28ybHydXOHihQOoQiIjcBntcqEGtgv3g5801XVxp3dEcXNdb73dERESWWLiQ3XgryHUe/18aXlt5ROowiIjcHgsXIjfx8/6LUodAROT2WLgQERGRbLBwISIiItlg4ULkRmb9dBD6+pbUJSLycCxcyG6Ojs1dNKWvS+JQsh/2XbRYYZeIiCyxcCG7hTbzNn/98tjODbYfkxCFv94eh8wF410ZluJcK6mweXzP2as4nVckcjRERO6FK4aR3dq1aI55E7shrJkPzl0pses5GjXnUDvKaGPrhPNXSnHvop0AwEKQiDwae1zIIZMT2+DWblEOP69n6yAXRKNM7649aXXsr/xiCSIhInI/iu1x0ev10Oudtzld1bWceU05M5qsV3ntEmW5O3TNXC17tD/2nruGB79Oc3lsSlRcVn37aFtGLtYezcVLo+Ph34RtFviZFgfzLA7mWRyuzLO911QJgjK29E1OTkZycjKMRiMyMjKwZMkS+Pv7Sx2WYqVmqfDLecttAD5KMmD6Ti+LxzUJAjBjl2JrZZdqHyDgryLL224jo024rQ1nIBGRMpSWlmLSpEkoLCxEYGBgne0UU7hU0el0CAoKQn5+fr1v3FF6vR4pKSkYNWoUtFqt064rV19sO4t3/zxlcezU3NGIn73O4nFt05amY92xPJfH5wnGJkRi4X09G/18fqbFwTyLg3kWhyvzrNPpEB4e3mDhoth//mq1Wpd8eF11XbnRqK03XaydF1t5Uqs4rMppVCpUmFSYtmQ/xnVriXv7xzTqMvxMi4N5FgfzLA5X5Nne6/GvCEnuhyeSrI75qBXVEegSggD8d+tZbDp5GbN+PiR1OEREomDhQpIb0DYUY2vNVIoPsl24rHnuJjFCkgUBAgrLOBCRiDwLCxdqsvsHxOLLqf2adI1PJ/fBu3f1aLBd12jnjVuSu9TjHCtERJ6HhQs12fw7u2NEl0gAwNLHBqJ9i2ZY9vhAh66hUqng6109bobL1jXMYOLtNCLyPCxcyKmS2och9YVhGNguzOHn1lWstG9huT7M0sccK4qUrNxgvZ4OAJzOK8bVkgoIgoAjWYUorTDYbEdEJDcsXMhtDI1vAQBIiA6wOD59ZEeLx0ntHS+KlOr73eetjm08mYeRH2xGn7kp+ONIDv728TbckbwdAHA2vwRD3tmA73adEztUIiKnYOFCbiPIX4sTc2/FiicGQuWke0XPDO/gnAvJyENf7zV/vWJ/5U7TGbmVWwbMXnUEF6+V4fVVRySJjYioqVi4kFvx1WqgrrUxY9eWAXW0btiLYzo1NSTZqOu2UZXV6VnYdjpfpGiIiFxDsQvQkXJ0iAjA8scHIiLQV+pQ3NrGE3m4tVvLOs9PX5Zu8fhIlg5Gju8lIplh4UJuqfadosRGDPb1NI5u3jHx810YHKnGba4Jh4jIJXiriNxSY4a43N4r2ulxyEl+cTne/OWoxbGGxgptz+WvACKSF/a4UKP4aq33KnIqByuX1iF+mH9nd6xOz3ZNPDIwe/XRhhsREckc/7lFjfL3/jFIbBuKf97a2bHnDWjcRoANiY9o7vCtEiIikh8WLtQovloNlj+RhKeGtXfoecM7RWDji8PQKbL+mUKO3iqac1sCNGqut1vbsWyd1CEQETkVCxcSXdvwZtB61V9kNHPgJubXD/VHXHgz+Go1eG1cF7ues/HFYfa/gIxlFZRJHQIRkVOxcCFJPHVz5cJwtXeFrjIuxoTB7cPw0X29HLruY0Pbwc+O8Texof7mrx9MauPQaxARkXRYuJAkxvdoiW3/HI5PJvWxeb6ZFlj8j764vVcrh6/97SMDEBXoi88f6AvAdnHEm0qWyiqMEOoZJJRfXI5vdmSisFQvYlRERNZYuJBkWof4N3pcSn29Kv3jQrHr1RG49UbB8sG9vaza1Jwm7OmDes9dKUWXN9Zi2pL9dbZ56Ou9mPPLUcz8IV28wIiIbGDhQrL054yhdrf1867/1lGQn7ap4cjadzc2alxzOKfONoezCgEAqSfyRImJiKguLFxIlmLD/BtuVA9VjS6X4Z0jmhoOERGJhIULeTxn7URNRESux8KFPM6orpFSh+BWvt9zQeoQiIjsxsKFZE/tQJdJz9ZB5tlGTbFoStOv4S70NbaIrm9mERGRO2DhQrI1KTEWfWKDMbi9/TtHx4Y1a/RMppZBvuavu7cKatQ13N2uM1elDoGIqF4sXEi23p7YHSueHgwvjf0f476xwY16rXXPD3WoZ0euvtmRKXUIRET14u7Q5BE2vjgMu89cwd19W1uda6gc6R0bjI4N7K2kFGuP1j0lmojIHbBwIY/QNrwZ2oY3s3muoVEdtp4XGeiL6CBfaL3UOHel1AkREhGRPXiriMhOT97YCXt895bQqFXYMms4NrwwzKpdfERzkSNzjR/2XcBPaRcbbJddUIZ31p5ANjd0JCIRsMeFPF5Dt4pUN1o8kBiLpHahiAur7IGpa2xNfGRznMordmaIotNd12PWT4cA1L0RZpWHvt6Lk7lFSD2ei3XP3yxGeETkwdjjQtQA4cbNJJVKhQ4RAQ0OBk5sa/8sJ3dkMJqwokZPS0m5od72J3OLAAAZufIu1ohIHtjjQh6vVYifU64TG+qPtyYk1DmWRi4S307FlZIK8+OlXKCOiNyIW/a4TJw4ESEhIbj77rulDoUULOX5oVjx9CBEBPg23NgOwzq1wPDOEbLfQqBm0QIA566WSBQJEZE1tyxcpk+fjm+//VbqMEjh4iMD0Cc2pMnXWT1tMB67qS1eGtPJCVG5nxX7s6yOmUwCDEZTnc/Zf/4a8ovLXRkWEXkotyxchg0bhoAAz1g3g9zThhfsH2TaMyYYr43vigBfrQsjch+CIODWj7bg5vc22Sxedp+5gjs/3YEB89ZLEB0RKZ3DhcuWLVtw2223ITo6GiqVCqtWrbJqk5ycjLi4OPj6+iIxMRF79uxxRqxEovjvg/3QroUypjS7QttX1iAjtxhZBWXIqjUF+ue0i/j7f3YBAEzc9oiIXMDhwqWkpAQ9e/ZEcnKyzfPLly/HzJkzMWfOHOzfvx89e/bEmDFjkJeXZ27Tq1cvdOvWzeq/7Ozsxr8Toib47pFEjO4aid2vjsDIWrtHB/g4NoY9JsQfnaM8o8ewwmDZ4/LCjwclioSIPIXDs4rGjh2LsWPH1nn+gw8+wGOPPYaHHnoIAPD555/j999/x1dffYWXX34ZAJCent64aG0oLy9HeXn1vXSdTgcA0Ov10Ov1Tnudqms585pkmxS5TowLQmJcT4vXXTAxASsOZGPasLYOx7L6qYF4/LsD2Hwq3+b5Di2a4fRl+Q96HfV/W+o9z58X/u4QC/MsDlfm2d5rOnU6dEVFBdLS0vDKK6+Yj6nVaowcORI7d+505kuZzZ8/H2+99ZbV8XXr1sHf39/pr5eSkuL0a5JtUufaD8DklsDOTY0bq3H5shp1dWqqK4rqPKcka9askToEtyH159lTMM/icEWeS0vt2z7FqYVLfn4+jEYjIiMtu9ojIyNx4sQJu68zcuRIHDx4ECUlJWjdujV+/PFHJCUl2Wz7yiuvYObMmebHOp0OMTExGD16NAIDAxv3RmzQ6/VISUnBqFGjoNV6xiBMqSgl1yuv7Mexguoelzah/jh3tfIHMzwsHBmFV6UKTTT9bhqBiAAfqcOQlFI+z+6OeRaHK/NcdcekIW65AN369fb/C9fHxwc+Pta/GLVarUs+vK66LlmTe65VNRZ06RkTjPv7x+DlFYcrz6llvtiLnaYuTsP6mdwGAJD/51kumGdxuCLP9l7PqX3V4eHh0Gg0yM3NtTiem5uLqKj69zshUrK+sSEN7kKtRKdlvmcTEbkfpxYu3t7e6Nu3L1JTU83HTCYTUlNT67zVQ+QJ1CpAqFG5qBrc2pGIiGxx+FZRcXExTp8+bX589uxZpKenIzQ0FLGxsZg5cyamTp2Kfv36YcCAAfjwww9RUlJinmUkFs4qki+l5NpUo1IRBBP0hurNCk2muledVZra38ej2TqkHM/DEze1hZ+3RqKoxKOUz7O7Y57F4Q6zilSCIDjUg71p0yYMHz7c6vjUqVOxePFiAMAnn3yC9957Dzk5OejVqxcWLlyIxMRER17GYcnJyUhOTobRaERGRgaWLFnikllFRPZadFyNYwWVnZoTYo0I8QG+OVX5h7pjkAkZhWrz13fFmTD/oFsOOWuyj5Isd5eevrPyfY6INmFCG88p4IiofqWlpZg0aRIKCwvrnVzjcOHi7nQ6HYKCgpCfn89ZRTKllFw/+r/92JxROavoyBsj4KVR47XVR9E3Nhi/Hc7Bjr8qZxWdmjsaRdf16DNvo5ThusyaZwchPqJ6JeL42esAADfHh+O/D/aRKizRKOXz7O6YZ3G4elZReHh4g4WLMv+JB84qUgK551pdY1ZRc//KHaj/fW9vAMCaI9UrSWu1WngZxY1NTOM+3oHMBeOtjqvVKll/fx0l98+zXDDP4lDMrCIiInupVKp6d5gGgMtF5Xh+eTr2Zip/zRsisg8LFyIXub1XKwBAuxbNrM51iLDcxNHRG7YBPl64o1d0o2MT28YTedjxVz5mLk83HzuWrUPn2WuxMPVUnc+b88sRrDyQhXs+d83K20QkP4q9VUQktdt7RSM2zN9ifEeVF8d0gkatwt96tGzUteU2MO2hxXutjuXorgMAPkjJwOiESHSOsr6nveZwjstjIyJ5UWzhwunQ8qWkXHdvWVm01H4vPmrg5THx5nMGB9+rAAEmk9zKl7qNX7gNJ94aZXEsq6DM4nFdn4fDWYX4/XAOpg1rjwBf9/uVpqTPsztjnsUhy+nQ7orToUnOSg3AK3vt/6ProxHQLURAWr5y7vbWnjZ9qRRYUGOKeO3zVaqmVw+ONOHedpxeTSRXnA7N6dCy5Ym5LizTo9/b9k+HbuajwYhOEfjl0CUXRiWurS8NRVSgr/nxqdxijPtkR/XjuaNtPq9qenW/NsFY+ugA1wbZCJ74eZYC8ywOTod2IU6Hlj9PyrXWdmdC3YTK6cRK8vYfGfjsgb7mx15ay19PDX0W1Cq1W39ePOnzLCXmWRxSTodWbOFCJCdBflqM6x4Fg1HA1EFxOJtfgtdXHZE6LFEVlnFsAhE1TDk3yIlk7tPJffGfB/thcIdwPDCwjc0279zVHVqNCp9P6WvzvJypanUgOXwTW1kdUERUBxYuRG5q9bTBGNqxBW6KDzcf+3v/WJyYOxY3xbeQMDLX4I7ZRGQPxd4q4nRo+WKuK3WNaoYvp/TGgrUnsfVU5Z5HVTkxGWExHXr6Le3x0Ya/JInTWbadzsfVojIs2XMBt/dqCYPB8vtf8/NQWmHAg1+nYVjH6qIOguCWnxl+nsXBPIuD06GdiNOhSanyyoB56V7oHGTCU12rp/v+75Qa+25Mh/4oyWCeFqwU/+xpwDs1pkMntjBhQIQJHQKBzZdUWJGpsWjfIdCEZxM4HZpIrjgdmtOhZYu5tlZ03YDmPhqoagwEefGnw1h9sHI69Km5o83TgpXi92eSMP4T66X+T80djS+2ncW7f1puFZDYNgTfPdxfrPDsxs+zOJhncXA6tAtxOrT8MdfVQm3koWYRo8Q8fbHtnM3jWq0WGrXG6rhK5d67TfPzLA7mWRzcHZqIqJZV6dkOtffEwb2Lt5/FM0v2N7jLNpGSsHAhkilV7fnDdujS0nm3T0l6b/56DL8duoQ/jnAzSvIcLFyIPMiCO7tLHYLL1FfHXSkux9trjuN0XrF4AYmopNzRpZeJ5IuFC5FMNfexHKK25NFEiSJxfy/9dAj/2XIG4xZulToUImoixQ7OJVK650d1xMmcItzdrzUAYFCHcJvturcKwuGsQihla6Oa69fYK/1CAQCgwsCxIERyp9jChQvQyRdzbZ8AbxW+f6QfgOpcfTapF77cnol95wrM7abf0g5/XS7BiM4RKFbALYVNJ3JgNBmtT9S7AF11sSP250qMz7PJZPT4nxf+3hAHF6BzIi5AR1SpsAJ4I6363yTPJRjQ/saYXEEA/nNCjWAfYEeuPO8UdwoyoXUzIDXbMv52AQLua29EpJ/1c17bq0GxobLL6aMk+RdvVaoWHbyvnRFJkYr4VU4ejAvQcQE62WKumyavqByD391sfrz00f7o1ybEqp3SFqyr8tXUPkhqGwqNWmWeeTVwwSZcKakAULmAnZhc+Xmu+h6+fUdX3NO3tVOvLTf8vSEOLkDnQlyATv6Y68YJalY9mCUm1A/92oZDq5Fn70pjLNlzES/+dARJ7cKQPLkPAMsZR1J9plz5edZoNPxZuYG/N8Qh5QJ0ii1ciDxVcx8vfPWPflCpVLg5vgXUTRyVe3PHFnj0praY8uUeJ0XoWuuP5wEAfj98CcnmowoZmVwHT1x8jzwXCxciBbqlc2STnh/e3Bv5xZW3Vhbe1xtB/lqLY55oX+ZVHM3W4cGkNo1a/I+InMNz+o+JyMIH93RHkFaAt5f1rwFNjV6aIP+q7lv5/bHe8Vc+9EZTvYvT2evuz3dizi9HseFEXtMvRkSNxh4XIg91W4+W0Fw8gHVFLfF7rSXjbQ3Zl2Mnw6QvduPO3q2ces2z+SVOvZ5TyPB7Q9RY7HEhIrPoIF+M7hoJW1MN5fq3ccWBLLhq8mSFwYQ3fzmKjeyFIRINCxciMtv6z1uwaEpfm+f+3j9G5Gicp66xORUGE/acvQq9g7srL9r8F7aeuowlu89h8Y5MPLR4r8MxZRWU4XRekcPPs8WeotJoElxWwBGJiYULEZlVrX1i6+/bcyPizV8rZfuA11cdxr2LduLNX47a/ZzNGZcx/48TmPLlHly8Vtbo1x68YANGfrAFV4rLG30Ne+mu69F/3no8u/SAy1+LyNUUO8aFS/7LF3Mtjqr8mgST1bGGlshXq1QwyfRf7xevFGHIe1ssjn2/+zweG9IGX20/h6lJsfjP1rPIvFKKb/7Rz2KgMgBsPZVv/tpgrN56oK7Pa3lFBYymus+fydMh0Ce4ke+mktFY/5L/vxy4iKslFfjt0CX83z3K/Lni7w1xuMOS/4opXGou+Q8A69atc8mS/ykpKU6/JtnGXIsjJycHVZ2va9asAQBcL9eg6gZE1bFKlb8yTIIJI6MFrM+WX6dt7aKlyn2fbkFOmQrf7jpvPvbpD3+Yt0uw9evy7NlM1M5dbe8d0uBauQYCUmA5gavyejt27EB2gGPvofY1Dh86BL+cg3W2OpyrAqCpN06l4O8Ncbgiz6WlpXa1U0zhMm3aNEybNs285P/o0aO55L9MMdfiqMpzVFQUcKVycOm4ceMAAJ+d3YkTOUUWxwDgw4xtOHulFMM6RuDTyb0wdXEadp65Kn7wLpBTZn3/K3HgQAyIC4XBaAJ2rrc63zq2DZBzAUDlvkEvjoqH7roeN3cMx4C40BvHK5flj04YgH5tq3fwrjo+aNAg9IoJblTMVdfo0bMHxtUze0q39yKWnzkGwPL7qST8vSEOVy/5bw/FFC61ccl/+WOuxaFWV3cDVOX708l98NrKI3jmlg4W34Oljyfhl4NZuLdfDLy9vbH08SS8u/YEPt30l+hxi8FL4wWtVoul+zJtnv9+zwWLx++nnAIA/GdrJjIXjLe8lpeXzc+zMz7nVXHWRaPRWLxeU5zI0eH8lVKMTohq0nVchb83xMEl/4nIrbRr0RxLHx9odTwqyBePD21vceyRIW2x6kAWsguvixWe6DadvOzwc/RGExZvz3R+MBK79cOtAICVTw9C71jrzTuJXE1+N6iJyKkmD6ic5jywXWijnh/W3AfbX77F4tiqaYMxNalNk2OT2t//swvD3tuIjFzHpy0v3p6JeWuOmx83ZihzfdOXa56TYnHAxuSEyBnY40Lk4frHhWD3qyMQ3tyn0deovXdPr5hgrDl8qamhuYXMK/YNGKzt4MWCRr9mWYURXd5YCwBY89xN6BptPV5PppO6iJqMPS5EhMhAX6tpv011U3x4w408yNK9F/HLwWxUGBpe7G7RluoxQ7N+tj1bqLF1y+r0rEY+05JJAP44fAnZBWW4rjfil4PZKCy1nM5aVmGs49lEjcfChYhc4qb4FlbHmnlrbLRUptJaf7RXHsjGc0sPIHnj6QZXsP1w/Snz18XXDTbbmxy4VSTUKHOmL0uvv3ENZRXGOlcVXrk/C099vx+DFmzAv347hueWHsBDi/eYz7+28jC6vLEWR7IK67z+tlP5SDunjFlpVYwmAc8vT8f/dmZKHYrZD3sv4NmlB+wqmuWAhQsROVWnyOpFSfa8OgIzR3U0P659S6m2u/q0dllcYqtrF+k/jlxCVkH1irtbMywH/tY8B1Teqnpo8V78nHYRF66W4vglHT5OPYUpX+42t3l++UGL3o5DFwvww94LKKswQhAEFF83WFzzp7SLePzbfUg9noudf13B0exC/HH4Er7bdQ7Tvt+PCoMJpRUGdJ2zFglz/sRDX+/B5lpx7smsLjhW7q/sxdl/vgA7/7qCOz/dju93V66H81HqKdhypbgcD3y5G3d9ttPurQhO5xVjdbrr9p4qrTDgh70XkG/nasZrj1zC0ezCWsdysPJAFmavtn81ZlcymQTM+vkQfj2YjZ/SLlqdzy8ux8OL9yLlWC6u640oN7h/LxnHuBCRU3w2uQ8+SMnAwvt7m49FBPriuRHx+CAlw3xsbLco/FFrN+oqRpMy/kVYn4zcYgx5Z6P58aGsQmQVlGHwgg2YN7Ebbupg3VO16eTlBmc29fzXOsSG+uPLqf0w4ZPtAIBZPx+y2fbFHytvP607lmvz/M0dW2DZ3vMQhMr9nDaevIyNJy9bTfGuUqav/mN3/xe7bLb582gOWof4ISE6CMXlBlwtqd4/yiQAaZmVe0b1iwvBrwcvYUiHcEQF+VpcY+QHmwEAPl4a3Nqtcjr2sWwdTubqMLZrhM3XtcfR7EK8u/YkrpZU4PCNHqKVTw/C/60/hdfHd0HHG8V4hcEEAQL0RgGncovw5Hf7AQBvTUjA1EFxAIBLhdWFp8kkQN3ALdiPU08hq6AM8+/sbi7syw1GqKCC940VC9cdzUF0sB+6tQpy+L21e7V6wUHd9eritrjcgOY+Xnh7zXFsOJFnUWj/Mf0mdGnpvHXQnI2FCxE5xdjuLTG2e0ub5+7u2xo/pV3EcyM64I7ereosXMZ2b4lV6dmuDNPtpBzLRcqNAuK1lUeadK3zV0sx6v9srwzsiLoKnoMXChy+liAIWL73PP7582EAwIyR8fhw/SkM61RdoF3XG3Hvop0AgGB/LQpu9B5N6BmNf9/bE1qNGpeLqntB1h/PRbnBiGB/b0z9qvL2VOntXVFUCsTPXoc+scH4+alB2HXmKu7/Yhdu6xmNO3pFY0SXSIvYfth3AZszLiPlWK7VbZSJn+4AAIzOuIxNLw5DqxA/dHz9D/P5B2vMmpvzy1EUlxtwraQC/9121ny88xtrcVOHcPSNC8GYhCi0b9HcnJNygwm+Wg3+faOov7tva1y8VoZB7cMw4oPNUKtU+PmpQZjy5W5curHUQFXhKAgCLl4rw5xfjpoLjjNvjzMXSafzinHmcjF2/HXF6vvx8s+HkFdUbn6et8b6xstdn+3AsX/disz8EhzKKsS4blHw0qjxx+FLiArwlnxguEpQ2HahVSvn5ufnc+VcmWKuxSFmno0mAWfyS9ChRTNcLq7A4Hc3W5yf0KMlnhvRHt4aNYa+3/Q/vOQ+tBoV9Mb6/8w8OiQO/92WafPcO3cm4M7erTDs31uQVVD3WkFD2odi21/Vt69+eToJEz7dadHm4UFtcLm4HPPvSICPVoP42evsfh/dogNxJNu+lV3rcuKtUdCoVXjsf/uxKSPf4lzV9YP8vFBYVnlrL7y5t8XO5qfmjgYATP16H3bUWrG6e6tAHM6qP764MH+7Z8mpVNUz14Z0CMOTQ9viga/2mc8fe2O4S1bODQ8PR2FhYb1/vxVTuNTcqygjIwNLlixxyV5FRNQ0pQbglb2Wnb3PdzMgLgAoKAfm7K8+l9jChO6hAkJ9BLx7iB3EnurVXga8ne7c7/+4GCPWXJDXYPF5/Qx4bZ97/Bx8lGRouJGDSktLMWnSJM8pXKqwx0X+mGtxSJnn73efh1qtwhu/VC7Q9sPjA9A7JhiXCq+be1z+3q81/t/tXc3PceRfx0TkWmumJSI+yvExN/Wxt8fFPUo3F+BeRfLHXItDijz/Y0jltgFVhYv2xj4+Xl7V/4p7/W9d64zr5bGdseCPE64PlIhsmv/nafzvUettQZrC3t9DnA5NRJKzZ9XeoR2rB3Pe1jPaleEQUQNqziQTm2J7XIjI/X33SCIKyioQE1o5Hq3mfevaa74kT+qN7m/Wfbsoc8F4ZBeUIdhfi+93nbfYJ4iIlIOFCxFJZkg92wLUXv3CS13dQazVWJ59eljlrafoYD8AwGND2+F0XjFaBPjgk42nLdrePyAWN8WH4+nv9zchciLPJuXoWBYuRCQLft4aPHtLB5QbTIgIqF6YbPnjA5HYLsyq/Tt39wAAnMjRYf3xyjUrxiREYtaYTghp5m3Rdmy3KOQVleP7RxOxJeMyHv9fmgvfCZH8mSSsXFi4EJHbaGiS4wujO1kda2gbgbcndkfmld24r38MHr2pndX5m+LD8dkDfc2PRydE4ffnhuBfvx7D7rPK2keHyFlMEva4cHAuEbmNAJ/qWQVaGyt62tJQsRMR6Iv1M2+2Klrm3p6AR4e0xbcPD7B6TkJ0EJY/kWRecr0uvz4zBP3jQuyKk0hJ2ONCRAQgyF+LRVP6QqtRNVg0VPH3btyvsSlJcQ22mT4iHu/9ebLO891bB+HHJwfhaHYhxi/cBgAIa+aNZ2/pgDd/PdaouLq2DIS3lxrp9Syvf0vnCPSJDcbjQ9ujpNyAub8fw4obGx0SiYGFCxHRDWMSouxq98bfuuLitTJ0a+W6zeCevLk9+sSGIKFVII5m6dAzJgi7zlzBtzvPYWqNwic+onpH7C2zhqOZjxd+O3QJ+85ds7rm2G5ReHxoO7QO8cd/t53Bos1nLM5//kBf5BZdxz2fVy9X3za8Gc7mlwAAnhjaDq+M62I+5+3ljYcHt3W4cMlcMB5xL/9ufhzWzBtXamx82DLI17xHjjPcPyAGS/dccNr1qtRcmp7EM8jGuDKxsHAhIll6eEhbl7+GRq1CUvvKX9BV/7+lcyRu6Wy5WZ+3lxrH/jUGapUKvtrKZeSXP5GELacu46Gv9wIA/nh2EFJO5GNqUhyC/CtvialqzJ06PW8svG7cHosN88e/bk+Av7cXsq6V4f7EGFwr0SPlWA4eGWI9Tqem9+7ugZd+OoQuLQOxetpgdJvzJyqMDe+6nTZ7FMYv3Iqj2Trc3isaH93X26KwaaoWAb5Y+tjAOnePbox1zw9Fnq4cD3y522nXtNfOV25B0vwNLn2NID8tCsv0DTd0wF19WuPn/RftavvOXd3Nm2PWFO4r4MVR8U6NyxEsXIiInKD2LSuNWoXWN6ZnA0CHiObo0spyPMyt3aLw+ea/EBnoYy5aqjxY61ZWRIAvOkUFwJYA3+rXvqdfDO7u29o8aDlj3lh0f/NPFF2vXJX4P1P6olsr20u1f/9oInb8dQUjukQAAPy0GquFxm7vFQ2DUYDBZML643kw3hil6afVYPOsYdhz9ipW7M9Cv7gQdIkKxEOL96JNmD8eH9oOzX28MK57FNYczsELozri3ykZ8PZSI8Rfi1xducXrvDSmEwJ9vTB79VH86/YE9GwdjHm/H8eAtqHmKe4dIwNQVlH3QmhPD2uPn/dfRK6uHLf3isaG43koKq97j531M4ciV1eOcoMRHVoE4NdD2Ui/UICC0grszazsPXswqQ0MJgEtg/zqvM7sv3XF3N+O4Ymh7bBoS2WP2oC2oRjVJRLz1hzHhJ7R+OVg5S7ofWKDseLpwdh+Oh+T/2tZgP38VBIMJgG3fri1ztfqFRNc723FEH8trt3YbXtwhzC8OaGruXD58O+9MGN5ep3PHd4pwvb7620070QtBRYuREQuEh8ZgGeHt0NO5imb53vFBCP1hZsRFehr87y92oQ1w8xRHRHkd6Mnp9ZMq0Htw/Dn0Vy0CPDB6Bq34t74W1f867djeHtidwBAsL83xnVvaT4/JiESq9KzEeDjhZSZN0OtqhzsXNMnG07h/XUZWHBXd0QE+OJvPaLxtx7VKxtnLhhv0f7TydUzuJ4dUfmv9iW7z+PVlYdrtOljjqPmWKQfnkwCANyfGAvfG2OgerQOwrTh7bHqQJZ59+hFU/pi08nLeOaWDnhxdCecyClCp6gAlFYYcLmoHDmF1zHpRpFw/4BYLN1zHn3bhKBDRAA61LjtN214BwDAR+tPmQuXf93ezSr/vWKC8c1DA7Bs73moVMAjQ9rikSFtYTQJKDeY0D8uFON7VL6fx4ZW9phVFS6TE9uYv0ePD22HzlEBGNk1EleKK9A2vBkA4Oz8cXh7zXHoygy4VlqBdcdyAQAzR3XEcyPicTqvCBeulqFNmD/2Zl616CU58MZoi56zAF+txfekvsKl5h24ri0DcexS03bHdhYWLkRELvTcLR2wZk1Gnefbt2junNcZUXfX/Tt39UC36HO4o3cri+MPD2mLu/u1RqCv7T1i3rq9G9q3aI4JvaIRFWS7uHrmlng8OCiuzmvYo3PL6mLh9fFdMLZb/eOcWtXoyVKpVHhpTGc8N6wdbl6wDu1ahmF010iLsVJdoyvHQQX4ahHgq7UY+P3quM54bXwX+Gsd3yn6l2cGY/vpK3jsprbw0qjxxM3tLc5r1Cq8OSHB5nNXPD0Ie89eNX9PVCoVXq0xdqlmPlUqFV4bX7nh6HW9EZ1nrwUATE6MBQCLgqtdi+ZQq1SYt+Y4vpza3+I1/bSO/cmPCPDB0scG4vglHVqH+LnN+kYsXIiIFC7Y39vcu1FbfQVHkJ+2zufZew179IkNwecP9EVcuD86RzVusLVarcLLPY0YN65vg2v7eNe4LadWqdDMp/4/hYF+ts/3aB2MHq2DHY4VqHzPfWIdn0rvq9XgrQkJqDCYEFbHHl+1bxf++56e+HLbWbw5oatV2+du6YCFG07j5o4tsDnjsvn4r88MgUpVOcYrqX0YBEHAx/f3RscW/jixd7PDcTuTYgsXvV4Pvd55g5qqruXMa5JtzLU4mGdxMM/2GdGpcvBzY/Ok1+uhUgEGg6HBwiXET4MpiTHw9lLDWy00+Jr39G6J7acu4+aOLdzi+zipf2Uvjb2xTOgRiQk9Im0+59nh7fCPpFiczC0yFy77X7sFAb5eVm1v7Vr5/k848NqOsPeaKqGh1ZtkIjk5GcnJyTAajcjIyMCSJUvg7+8vdVhERESysO+yChF+AmKdc/fSYaWlpZg0aRIKCwsRGFh3z5tiCpcqOp0OQUFByM/Pr/eNO0qv1yMlJQWjRo2CVtu0blGqH3MtDuZZHMyzOJhncbgyzzqdDuHh4Q0WLoq9VaTVal3y4XXVdckacy0O5lkczLM4mGdxuCLP9l6PexURERGRbLBwISIiItlg4UJERESywcKFiIiIZIOFCxEREckGCxciIiKSDRYuREREJBssXIiIiEg2WLgQERGRbLBwISIiItlg4UJERESywcKFiIiIZENxmyxWbXat0+mcel29Xo/S0lLodDpu4OVizLU4mGdxMM/iYJ7F4co8V/3drvo7XhfFFS5FRUUAgJiYGIkjISIiIkcVFRUhKCiozvMqoaHSRmZMJhOys7MREBAAlUrltOvqdDrExMTgwoULCAwMdNp1yRpzLQ7mWRzMsziYZ3G4Ms+CIKCoqAjR0dFQq+seyaK4Hhe1Wo3WrVu77PqBgYH8oRAJcy0O5lkczLM4mGdxuCrP9fW0VOHgXCIiIpINFi5EREQkGyxc7OTj44M5c+bAx8dH6lAUj7kWB/MsDuZZHMyzONwhz4obnEtERETKxR4XIiIikg0WLkRERCQbLFyIiIhINli4EBERkWywcLFTcnIy4uLi4Ovri8TEROzZs0fqkNzW/Pnz0b9/fwQEBCAiIgJ33HEHTp48adHm+vXrmDZtGsLCwtC8eXPcddddyM3NtWhz/vx5jB8/Hv7+/oiIiMBLL70Eg8Fg0WbTpk3o06cPfHx80KFDByxevNjVb89tLViwACqVCjNmzDAfY56dIysrCw888ADCwsLg5+eH7t27Y9++febzgiDgjTfeQMuWLeHn54eRI0fi1KlTFte4evUqJk+ejMDAQAQHB+ORRx5BcXGxRZtDhw7hpptugq+vL2JiYvDuu++K8v7chdFoxOzZs9G2bVv4+fmhffv2mDt3rsXeNcy147Zs2YLbbrsN0dHRUKlUWLVqlcV5MXP6448/onPnzvD19UX37t2xZs0ax9+QQA1atmyZ4O3tLXz11VfC0aNHhccee0wIDg4WcnNzpQ7NLY0ZM0b4+uuvhSNHjgjp6enCuHHjhNjYWKG4uNjc5sknnxRiYmKE1NRUYd++fcLAgQOFQYMGmc8bDAahW7duwsiRI4UDBw4Ia9asEcLDw4VXXnnF3ObMmTOCv7+/MHPmTOHYsWPCxx9/LGg0GmHt2rWivl93sGfPHiEuLk7o0aOHMH36dPNx5rnprl69KrRp00b4xz/+IezevVs4c+aM8OeffwqnT582t1mwYIEQFBQkrFq1Sjh48KAwYcIEoW3btkJZWZm5za233ir07NlT2LVrl7B161ahQ4cOwv33328+X1hYKERGRgqTJ08Wjhw5IixdulTw8/MTFi1aJOr7ldK8efOEsLAw4bfffhPOnj0r/Pjjj0Lz5s2Fjz76yNyGuXbcmjVrhNdee01YsWKFAEBYuXKlxXmxcrp9+3ZBo9EI7777rnDs2DHh9ddfF7RarXD48GGH3g8LFzsMGDBAmDZtmvmx0WgUoqOjhfnz50sYlXzk5eUJAITNmzcLgiAIBQUFglarFX788Udzm+PHjwsAhJ07dwqCUPmDplarhZycHHObzz77TAgMDBTKy8sFQRCEWbNmCQkJCRav9fe//10YM2aMq9+SWykqKhLi4+OFlJQU4eabbzYXLsyzc/zzn/8UhgwZUud5k8kkREVFCe+99575WEFBgeDj4yMsXbpUEARBOHbsmABA2Lt3r7nNH3/8IahUKiErK0sQBEH49NNPhZCQEHPeq167U6dOzn5Lbmv8+PHCww8/bHHszjvvFCZPniwIAnPtDLULFzFzeu+99wrjx4+3iCcxMVF44oknHHoPvFXUgIqKCqSlpWHkyJHmY2q1GiNHjsTOnTsljEw+CgsLAQChoaEAgLS0NOj1eoucdu7cGbGxseac7ty5E927d0dkZKS5zZgxY6DT6XD06FFzm5rXqGrjad+XadOmYfz48Va5YJ6d45dffkG/fv1wzz33ICIiAr1798YXX3xhPn/27Fnk5ORY5CgoKAiJiYkWeQ4ODka/fv3MbUaOHAm1Wo3du3eb2wwdOhTe3t7mNmPGjMHJkydx7do1V79NtzBo0CCkpqYiIyMDAHDw4EFs27YNY8eOBcBcu4KYOXXW7xIWLg3Iz8+H0Wi0+MUOAJGRkcjJyZEoKvkwmUyYMWMGBg8ejG7dugEAcnJy4O3tjeDgYIu2NXOak5NjM+dV5+pro9PpUFZW5oq343aWLVuG/fv3Y/78+VbnmGfnOHPmDD777DPEx8fjzz//xFNPPYXnnnsO33zzDYDqPNX3OyInJwcREREW5728vBAaGurQ90LpXn75Zdx3333o3LkztFotevfujRkzZmDy5MkAmGtXEDOndbVxNOeK2x2a3Mu0adNw5MgRbNu2TepQFOfChQuYPn06UlJS4OvrK3U4imUymdCvXz+8/fbbAIDevXvjyJEj+PzzzzF16lSJo1OWH374Ad9//z2WLFmChIQEpKenY8aMGYiOjmauyYw9Lg0IDw+HRqOxmomRm5uLqKgoiaKSh2eeeQa//fYbNm7ciNatW5uPR0VFoaKiAgUFBRbta+Y0KirKZs6rztXXJjAwEH5+fs5+O24nLS0NeXl56NOnD7y8vODl5YXNmzdj4cKF8PLyQmRkJPPsBC1btkTXrl0tjnXp0gXnz58HUJ2n+n5HREVFIS8vz+K8wWDA1atXHfpeKN1LL71k7nXp3r07pkyZgueff97co8hcO5+YOa2rjaM5Z+HSAG9vb/Tt2xepqanmYyaTCampqUhKSpIwMvclCAKeeeYZrFy5Ehs2bEDbtm0tzvft2xdardYipydPnsT58+fNOU1KSsLhw4ctflhSUlIQGBho/iOSlJRkcY2qNp7yfRkxYgQOHz6M9PR083/9+vXD5MmTzV8zz003ePBgq+n8GRkZaNOmDQCgbdu2iIqKssiRTqfD7t27LfJcUFCAtLQ0c5sNGzbAZDIhMTHR3GbLli3Q6/XmNikpKejUqRNCQkJc9v7cSWlpKdRqyz9LGo0GJpMJAHPtCmLm1Gm/Sxwayuuhli1bJvj4+AiLFy8Wjh07Jjz++ONCcHCwxUwMqvbUU08JQUFBwqZNm4RLly6Z/ystLTW3efLJJ4XY2Fhhw4YNwr59+4SkpCQhKSnJfL5qmu7o0aOF9PR0Ye3atUKLFi1sTtN96aWXhOPHjwvJyckeNU3XlpqzigSBeXaGPXv2CF5eXsK8efOEU6dOCd9//73g7+8vfPfdd+Y2CxYsEIKDg4XVq1cLhw4dEm6//Xab00l79+4t7N69W9i2bZsQHx9vMZ20oKBAiIyMFKZMmSIcOXJEWLZsmeDv76/YKbq2TJ06VWjVqpV5OvSKFSuE8PBwYdasWeY2zLXjioqKhAMHDggHDhwQAAgffPCBcODAAeHcuXOCIIiX0+3btwteXl7C+++/Lxw/flyYM2cOp0O70scffyzExsYK3t7ewoABA4Rdu3ZJHZLbAmDzv6+//trcpqysTHj66aeFkJAQwd/fX5g4caJw6dIli+tkZmYKY8eOFfz8/ITw8HDhhRdeEPR6vUWbjRs3Cr169RK8vb2Fdu3aWbyGJ6pduDDPzvHrr78K3bp1E3x8fITOnTsL//nPfyzOm0wmYfbs2UJkZKTg4+MjjBgxQjh58qRFmytXrgj333+/0Lx5cyEwMFB46KGHhKKiIos2Bw8eFIYMGSL4+PgIrVq1EhYsWODy9+ZOdDqdMH36dCE2Nlbw9fUV2rVrJ7z22msWU2yZa8dt3LjR5u/kqVOnCoIgbk5/+OEHoWPHjoK3t7eQkJAg/P777w6/H5Ug1FiSkIiIiMiNcYwLERERyQYLFyIiIpINFi5EREQkGyxciIiISDZYuBAREZFssHAhIiIi2WDhQkRERLLBwoWIiIhkg4ULERERyQYLFyIiIpINFi5EREQkGyxciIiISDb+P6rSDyJZZFbOAAAAAElFTkSuQmCC\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "rng = nnx.Rngs(0)\n",
        "temperature = 0.5\n",
        "prompt = \"ROMEO.\"\n",
        "prompt_tokens = tokenizer(prompt, return_tensors=\"jax\").input_ids\n",
        "logits, state = prefill(parallel_model, prompt_tokens)\n",
        "sequential_model = to_sequential(parallel_model)\n",
        "token = jax.random.categorical(rng(), logits / temperature)\n",
        "streamer = TextStreamer(tokenizer, skip_special_tokens=True)\n",
        "streamer.put(prompt_tokens)\n",
        "streamer.put(token)\n",
        "for step in range(512):\n",
        "    token, state = sample(sequential_model, token, state, temperature=temperature, rng=rng())\n",
        "    streamer.put(token)\n",
        "streamer.end()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E1yxm5LmiXBZ",
        "outputId": "3d707a84-5700-4a82-a04d-bcf0df7ebeb9"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "ROMEO.\n",
            "A torch for me: let wantons, light of heart,\n",
            "Tickle the senseless rushes with their heels;\n",
            "For I am proverb’d with a grandsire phrase,\n",
            "I’ll be a candle-holder and look on,\n",
            "The game was ne’er so fair, and I am done.\n",
            "\n",
            "MERCUTIO.\n",
            "Tut, dun’s the mouse, the constable’s own word:\n",
            "If thou art dun, we’ll draw thee from the mire\n",
            "Or save your reverence love, wherein thou stickest\n",
            "Up to the ears. Come, we burn daylight, ho.\n",
            "\n",
            "ROMEO.\n",
            "Nay, that’s not so.\n",
            "\n",
            "MERCUTIO.\n",
            "I mean sir, in delay\n",
            "We waste our lights in vain, light lights by day.\n",
            "Take our good meaning, for our judgment sits\n",
            "Five times in that ere once in our five wits.\n",
            "\n",
            "ROMEO.\n",
            "And we mean well in going to this mask;\n",
            "But ’tis no wit to go.\n",
            "\n",
            "MERCUTIO.\n",
            "Why, may one ask?\n",
            "\n",
            "ROMEO.\n",
            "I dreamt a dream tonight.\n",
            "\n",
            "MERCUTIO.\n",
            "And so did I.\n",
            "\n",
            "ROMEO.\n",
            "Well what was yours?\n",
            "\n",
            "MERCUTIO.\n",
            "That dreamers often lie.\n",
            "\n",
            "ROMEO.\n",
            "In bed asleep, while they do dream things true.\n",
            "\n",
            "MERCUTIO.\n",
            "O, then, I see Queen Mab hath been with you.\n",
            "She is the fairies’ midwife, and she comes\n",
            "In shape no bigger than an agate-stone\n",
            "On the fore-finger of an alderman,\n",
            "Drawn with a team of little atomies\n",
            "Over men’s noses as they lie asleep:\n",
            "Her waggon-spokes made of long spinners’ legs;\n",
            "The cover, of the wings of grasshoppers;\n",
            "Her traces, of the smallest spider’s web;\n",
            "The collars, of the moonshine’s watery beams;\n",
            "Her whip of cricket�\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "q689f1q2hwEC"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}