fizzbuzz in jax

fizzbuzz.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3cc970e3-a4e2-40ae-a4da-db15abcd550c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from jax.tree_util import register_pytree_node_class\n",
    "import optax\n",
    "from tqdm.notebook import trange\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "jax.config.update('jax_debug_nans', True)\n",
    "\n",
    "model_dtype = jnp.float32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7b1f6362-de4b-4e09-84dc-7ccff11cd2be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a layer\n",
    "@register_pytree_node_class\n",
    "class Layer:\n",
    "    weights: jax.Array\n",
    "    biases: jax.Array\n",
    "\n",
    "    def __init__(self, weights: jax.Array, biases: jax.Array):\n",
    "        # weights array: columns are nodes, rows are weights\n",
    "        self.weights = weights\n",
    "        self.biases = biases\n",
    "\n",
    "    @staticmethod\n",
    "    def init_size(randkey, inputs: int, nodes: int):\n",
    "        weights = jnp.sqrt(2 / inputs) * jax.random.normal(randkey, shape=(inputs, nodes))\n",
    "        biases = jnp.zeros(nodes)\n",
    "        return Layer(jnp.asarray(weights, dtype=model_dtype), jnp.asarray(biases, dtype=model_dtype))\n",
    "    \n",
    "    def forward(self, input: jax.Array) -> jax.Array:\n",
    "        return input @ self.weights + self.biases\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'<Layer weights={repr(self.weights)} biases={repr(self.biases)}>'\n",
    "    \n",
    "    def tree_flatten(self):\n",
    "        return ((self.weights, self.biases), None)\n",
    "    \n",
    "    @classmethod\n",
    "    def tree_unflatten(cls, aux_data, children):\n",
    "        return cls(children[0], children[1])\n",
    "    \n",
    "@register_pytree_node_class\n",
    "class Network:\n",
    "    layers: list[Layer]\n",
    "\n",
    "    def __init__(self, layers):\n",
    "        self.layers = layers\n",
    "\n",
    "    @staticmethod\n",
    "    def new(randkey):\n",
    "        key1, key2, key3 = jax.random.split(randkey, 3)\n",
    "        return Network([\n",
    "            Layer.init_size(key1, 32, 1024),\n",
    "            Layer.init_size(key2, 1024, 512),\n",
    "            Layer.init_size(key3, 512, 3),\n",
    "        ])\n",
    "\n",
    "    def forward(self, value: jax.Array) -> jax.Array:\n",
    "        layers = iter(self.layers)\n",
    "        value = jax.nn.tanh(next(layers).forward(value))\n",
    "        value = jax.nn.tanh(next(layers).forward(value))\n",
    "        value = jax.nn.sigmoid(next(layers).forward(value))\n",
    "\n",
    "        return value\n",
    "    \n",
    "    def tree_flatten(self):\n",
    "        return (self.layers, None)\n",
    "    \n",
    "    @classmethod\n",
    "    def tree_unflatten(cls, aux_data, children):\n",
    "        return cls(children)\n",
    "\n",
    "def log_loss(out: jax.Array, expected: jax.Array):\n",
    "    eps = 1e-7\n",
    "    out = jnp.clip(out, eps, 1 - eps)\n",
    "    errors = -expected * jnp.log(out) - (1 - expected) * jnp.log(1 - out)\n",
    "    # sum horizontally (each data row)\n",
    "    return jnp.mean(jnp.sum(errors, axis=1))\n",
    "\n",
    "def loss(model: Network, x, y):\n",
    "    y_hat = model.forward(x)\n",
    "    return log_loss(y_hat, y)\n",
    "\n",
    "loss_grads = jax.value_and_grad(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cb24b1e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1695609639.709761  125941 tfrt_cpu_pjrt_client.cc:349] TfrtCpuClient created.\n"
     ]
    }
   ],
   "source": [
    "bit_width = 1\n",
    "shift_by = jnp.array([i for i in range(32 // bit_width)], dtype=jnp.uint32)\n",
    "mask = (1 << bit_width) - 1\n",
    "\n",
    "def split_nums(nums: jax.Array):\n",
    "    \"split u32 into bits\"\n",
    "    inputs = jnp.reshape(nums, (-1, 1))\n",
    "    return (jnp.bitwise_and(jnp.right_shift(inputs, shift_by), mask)).astype(model_dtype)\n",
    "\n",
    "def make_data_set(randkey: jax.random.PRNGKey, min: int, max: int, count: int | None) -> (jax.Array, jax.Array):\n",
    "    minval = jnp.array(min, dtype=jnp.uint32)\n",
    "    maxval = jnp.array(max, dtype=jnp.uint32)\n",
    "    if count is not None:\n",
    "        nums = jax.random.randint(randkey, (count,), minval, maxval, dtype=jnp.uint32)\n",
    "    else:\n",
    "        nums = jnp.arange(min, max, dtype=jnp.uint32)\n",
    "    inputs = split_nums(nums)\n",
    "    out_fizz = jnp.equal(jnp.mod(nums, 3), 0)\n",
    "    out_buzz = jnp.equal(jnp.mod(nums, 5), 0)\n",
    "    out_none = jnp.logical_and(jnp.logical_not(out_fizz), jnp.logical_not(out_buzz))\n",
    "    outputs = jnp.stack((out_fizz, out_buzz, out_none), axis=1).astype(model_dtype)\n",
    "    return inputs, outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f7b3ff70",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_step(optimizer, batch_size, min, max):\n",
    "    @jax.jit\n",
    "    def step(model, randkey, opt_state):\n",
    "        key, rand_split = jax.random.split(randkey)\n",
    "        train_in, train_out = make_data_set(key, min, max, batch_size)\n",
    "        loss_v, grads = loss_grads(model, train_in, train_out)\n",
    "        updates, opt_state = optimizer.update(grads, opt_state, model)\n",
    "        #updates = jax.tree_util.tree_map(lambda x: x * -0.1, grads)\n",
    "        model = optax.apply_updates(model, updates)\n",
    "        return rand_split, model, opt_state, loss_v\n",
    "    \n",
    "    return step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e64deb5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def model_eval(model, input):\n",
    "    return model.forward(input)\n",
    "\n",
    "test_threshold = 0.75\n",
    "def test_fizzbuzz(model, input, quiet=False):\n",
    "    out = model_eval(model, split_nums(jnp.array(input)))\n",
    "    for i, val in zip(input, out):\n",
    "        fizz, buzz, neither = val\n",
    "        fizzmod = i % 3\n",
    "        buzzmod = i % 5\n",
    "        if not quiet:\n",
    "            print(\n",
    "                f'{i}  fizz: {fizz:.4f}, fizzmod: {fizzmod}, ' +\n",
    "                    f'buzz: {buzz:.4f}, buzzmod: {buzzmod}, ' +\n",
    "                    f'neither: {neither:.4f}')\n",
    "        if fizzmod == 0 and fizz < test_threshold:\n",
    "            print(i, 'fizz wrong: expect fizz got', fizz)\n",
    "        if fizz >= test_threshold and fizzmod != 0:\n",
    "            print(i, 'fizz wrong: expect not fizz got', fizz)\n",
    "        if buzzmod == 0 and buzz < test_threshold:\n",
    "            print(i, 'buzz wrong: expect buzz got', buzz)\n",
    "        if buzz >= test_threshold and buzzmod != 0:\n",
    "            print(i, 'buzz wrong: expect not buzz got', buzz)\n",
    "        if fizzmod != 0 and buzzmod != 0 and neither < test_threshold:\n",
    "            print(i, 'neither wrong: expect neither got', neither)\n",
    "        if (fizzmod == 0 or buzzmod == 0) and neither >= test_threshold:\n",
    "            print(i, 'neither wrong: expect not neither got', neither)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "178e7313",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_stage(randkey, model, optimizer, batch_size, iterations, min_val, max_val):\n",
    "    step = make_step(optimizer, batch_size, min_val, max_val)\n",
    "    progress = trange(iterations)\n",
    "    lossplot_v = np.zeros(len(progress), dtype=model_dtype)\n",
    "    opt_state = optimizer.init(model)\n",
    "    for i in progress:\n",
    "        randkey, model, opt_state, loss_v = step(model, randkey, opt_state)\n",
    "        progress.set_postfix({ 'loss': loss_v })\n",
    "        lossplot_v[i] = loss_v\n",
    "        if loss_v < 0.005:\n",
    "            break\n",
    "\n",
    "    lossplot_t = np.arange(progress.n)\n",
    "    plt.plot(lossplot_t, lossplot_v[:progress.n])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b2df9625",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Network.new(jax.random.PRNGKey(42))\n",
    "optimizer = optax.adam(learning_rate=0.0005)\n",
    "batch_size = 16384\n",
    "randkey = jax.random.PRNGKey(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ac8334dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c51f20058c174440ae27fc3e778fac62",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16384 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0  fizz: 0.9997, fizzmod: 0, buzz: 0.9998, buzzmod: 0, neither: 0.0000\n",
      "1  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 1, neither: 1.0000\n",
      "2  fizz: 0.0000, fizzmod: 2, buzz: 0.0001, buzzmod: 2, neither: 1.0000\n",
      "3  fizz: 0.9999, fizzmod: 0, buzz: 0.0002, buzzmod: 3, neither: 0.0000\n",
      "4  fizz: 0.0001, fizzmod: 1, buzz: 0.0000, buzzmod: 4, neither: 1.0000\n",
      "5  fizz: 0.0090, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "6  fizz: 0.9999, fizzmod: 0, buzz: 0.0002, buzzmod: 1, neither: 0.0000\n",
      "7  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 2, neither: 1.0000\n",
      "8  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "9  fizz: 0.9999, fizzmod: 0, buzz: 0.0001, buzzmod: 4, neither: 0.0000\n",
      "10  fizz: 0.0056, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "11  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 1, neither: 1.0000\n",
      "12  fizz: 0.9999, fizzmod: 0, buzz: 0.0001, buzzmod: 2, neither: 0.0001\n",
      "13  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "14  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 4, neither: 1.0000\n",
      "15  fizz: 0.9999, fizzmod: 0, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "16  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 1, neither: 1.0000\n",
      "17  fizz: 0.0090, fizzmod: 2, buzz: 0.0004, buzzmod: 2, neither: 0.9880\n",
      "18  fizz: 0.9999, fizzmod: 0, buzz: 0.0003, buzzmod: 3, neither: 0.0000\n",
      "19  fizz: 0.0001, fizzmod: 1, buzz: 0.0016, buzzmod: 4, neither: 0.9986\n",
      "20  fizz: 0.0106, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "21  fizz: 0.9729, fizzmod: 0, buzz: 0.0000, buzzmod: 1, neither: 0.0226\n",
      "22  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 2, neither: 0.9999\n",
      "23  fizz: 0.0006, fizzmod: 2, buzz: 0.0000, buzzmod: 3, neither: 0.9997\n",
      "24  fizz: 0.9999, fizzmod: 0, buzz: 0.0001, buzzmod: 4, neither: 0.0000\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAAGdCAYAAADAAnMpAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA150lEQVR4nO3deXxU5d3///eZJDNJIAsQskHYEWRHEIxWwBIF5LZS+/XGpcXaam8t3q3FWyuutVts/Wntt6Lo7c96d1GsVbC3WiqyqgQUJCwKQdawJGHPhJB9ru8fyIGRBBJIcmVmXs/HYx5eZ5v5XDmj8/ac65zjGGOMAAAALPHYLgAAAEQ2wggAALCKMAIAAKwijAAAAKsIIwAAwCrCCAAAsIowAgAArCKMAAAAq6JtF9AYgUBAe/fuVUJCghzHsV0OAABoBGOMysrKlJmZKY+n4eMfIRFG9u7dq6ysLNtlAACAc7Br1y517dq1weUhEUYSEhIkHe9MYmKi5WoAAEBj+P1+ZWVlub/jDQmJMHLi1ExiYiJhBACAEHO2IRYMYAUAAFYRRgAAgFWEEQAAYBVhBAAAWEUYAQAAVhFGAACAVYQRAABgFWEEAABYRRgBAABWEUYAAIBVhBEAAGAVYQQAAFgVEg/KaykvfrBNuw9X6IZRWeqfzgP4AACwIaKPjLyzvkgvL9+hwoPHbJcCAEDEiugwcoKxXQAAABEsosOIY7sAAAAQ2WHkBMOhEQAArInoMOI4HBsBAMC2iA4jJ3FoBAAAWyI6jHBcBAAA+yI7jHyZRhgzAgCAPREdRgAAgH0RHUacL0/UcGAEAAB7IjqMAAAA+yI7jDBmBAAA6yI7jHzJcKIGAABrIjqMcGkvAAD2RXQYOYHTNAAA2BPRYYS7wQMAYF9Eh5ETODACAIA9ER1GHEaNAABgXUSHkRMMg0YAALAmosMIY0YAALCPMAIAAKyK6DACAADsi+gw4j4ojyEjAABYE9FhBAAA2BfRYeTEmBGeTQMAgD0RHUYAAIB9hBExZgQAAJsIIwAAwKqIDiOOw9U0AADYFtFhBAAA2BfRYaSktFKS9MEX+y1XAgBA5IroMFJQUiZJmpe/13IlAABErogOIwAAwD7CCAAAsIowAgAArCKMAAAAqwgjAADAKsIIAACwijACAACsIowAAACrCCMAAMAqwggAALCKMAIAAKwijAAAAKsIIwAAwCrCCAAAsIowAgAArCKMAAAAqwgjAADAqiaFkdzcXF188cVKSEhQamqqpkyZooKCgrNu9/rrr6t///6KjY3V4MGD9e67755zwQAAILw0KYwsXbpU06dP14oVK7RgwQLV1NToqquuUnl5eYPbLF++XDfeeKO+//3va82aNZoyZYqmTJmiDRs2nHfxAAAg9DnGGHOuG+/fv1+pqalaunSpxowZU+86U6dOVXl5ud5++2133iWXXKJhw4Zp9uzZjfocv9+vpKQklZaWKjEx8VzLPU2P+99x2zsen9xs7wsAABr/+31eY0ZKS0slSR07dmxwnby8POXk5ATNmzBhgvLy8hrcpqqqSn6/P+gFAADC0zmHkUAgoLvvvluXXXaZBg0a1OB6xcXFSktLC5qXlpam4uLiBrfJzc1VUlKS+8rKyjrXMgEAQBt3zmFk+vTp2rBhg+bMmdOc9UiSZs6cqdLSUve1a9euZv8MAADQNkSfy0Z33XWX3n77bS1btkxdu3Y947rp6ekqKSkJmldSUqL09PQGt/H5fPL5fOdS2jk7eLRKndq37mcCAIAmHhkxxuiuu+7S3LlztWjRIvXs2fOs22RnZ2vhwoVB8xYsWKDs7OymVdrCHvnHZ7ZLAAAgIjXpyMj06dP1yiuv6K233lJCQoI77iMpKUlxcXGSpGnTpqlLly7Kzc2VJP34xz/W2LFj9eSTT2ry5MmaM2eOVq1apRdeeKGZu3J+NheX2S4BAICI1KQjI88995xKS0s1btw4ZWRkuK/XXnvNXaewsFBFRUXu9KWXXqpXXnlFL7zwgoYOHaq///3vmjdv3hkHvdoQOPcrnAEAwHlo0pGRxtySZMmSJafNu/7663X99dc35aNaRc6FqXp/4z5JElkEAAA7IvrZNP/3xuFuu440AgCAFREdRmKiTna/LkAYAQDAhogOI1GO47Y5MAIAgB0RHUY8npNhhAGsAADYEdFh5FScpgEAwA7CyJf2lVU16mohAADQvAgjp/BX1NouAQCAiEMYOcWhY9W2SwAAIOIQRk5xrJojIwAAtDbCyCmOVdfZLgEAgIhDGDnFn/N22i4BAICIE/FhpHuneLf9j7V7VVsXsFgNAACRJ+LDSM6FaUHTfR78p95YvdtSNQAARJ6IDyP3XHXB6fNeX6uPtx8643abiv1av7tUkrT9QLn+Z/kOVdUy5gQAgKZyTAjc6cvv9yspKUmlpaVKTExs9vdfvfOwvvXc8tPmd07w6dXbL1H3TvH6/ftf6JnFW3RhRqL++N2LdUnuQklSSnufDhytkiR1SY7TniMV7vbL7r1C3TrF67VPCjXzzfUKGOneCf101YA09U1LaPZ+AADQljT295sw8qW+D76rmrrW+1P0T0/QtOweivdGadehY7r5ku7q2M7bap8PAEBLI4w0UVVtnfo9NL9F3ruxPntsggLGKCE2xmodAAA0h8b+fkf8mJETfNFRyuoYZ7WGgY/+S4N/9p4WF+yTv7JGFdz3BAAQAaJtF9CW/HRif931yppGrdutY7wKDx1rkTpu/eMnbvvCjETlXJiqe67q1yKfBQCAbYSRU0wenKH+MxKUnhSn+JgoeTxO0PK6gNG/PivWsKxkZSafPIqy40C5NhX7FeeN1vrdRzT9ij5ynOPbLvi8RI++tUFF/krNnNRft1/eSx9vP6SpL6xoVE0bi/zaWOTXniMV2n6gXD8a31cju3fgVA4AIGwwZsSSwoPH5DjSvDV7tOPgMXmjHb368a5Gb//4dYN15YA0tY+Nli86qgUrBQDg3DCANQTtOnRMl/92cZO3m35Fb/37yCx179SuBaoCAODcMIA1BGV1jNeOxyfrznG91d7X+DNosxZv1dgnlrRcYQAAtCDGjLRBP53YXz+d2F/GGNXUGS3fekDfPWVQa0O27CtTZnKc4r3sVgBA6OA0TYioqq3T6h2H9eSCzVq98/AZ1137yFVKimeAKwDALsaMhLHq2oD8lTWa+PQH7q3oTzXmgs760/dGWagMAICTGDMSxrzRHqW092nVQzl6aPKFpy1ftnm/+xA/AADaOsJIiLvt8l76tyEZp82/5pkPtWzzfgsVAQDQNISRMPDMTRdp7aNXnTZ/2ksfa+W2gxYqAgCg8QgjYSIpLkYbHptw2vypL6zQ+5+XKASGBgEAIhQDWMPMmZ4+/PX+qXrpuxe3ckUAgEjFANYI5YuO0oqZ4+tdtmjTPn2y41ArVwQAwJkRRsJQelKstudeXe+y62fnyV9Z08oVAQDQMMJImHIcR2/cma2U9t7Tlg352Xtau+tI6xcFAEA9CCNhbET3jlr10JX6ziXdT1t27ayP1OP+d1RZU2ehMgAATiKMRIBfTBmkN+7MrndZ/4fnq6qWQAIAsIcwEiFGdO+opfeOU9/U9qct6/fQfP363Y1c/gsAsIIwEkG6d2qnBTPG1rvshWXb1HPmu61cEQAAhJGI9PGD4/X7G4apX1rCacuWFOzTB1/s1+7DxyxUBgCIRNz0LIIZY7TrUIXGPLG43uU7Hp/cyhUBAMIJNz3DWTmOo26d4m2XAQCIcIQRaO4PL613fom/spUrAQBEIsIINLxbB+14fLJm3XRR0PzRv15oqSIAQCQhjMA1eUiGFt4TfLXNg3PXW6oGABApCCMI0rtzez00+UJ3+q8rC3WovNpiRQCAcEcYwWkGZASPeH7iX5ssVQIAiASEEZxmeLcOQdOvfrzLUiUAgEhAGMFp4rxRWv+zq4Lmrdx20FI1AIBwRxhBvRJiY/Tpw1e601NfWKFAoM3fHw8AEIIII2hQx3beoOleD7yr2rqApWoAAOGKMIIz+tt/ZAdNbywqs1QJACBcEUZwRhf3CB7M+sGW/ZYqAQCEK8IIzshxHL04baQ7/dv5BRarAQCEI8IIzipnQJp6pbRzp78o4VQNAKD5EEbQKJOHZLjtv3+622IlAIBwQxhBo9x2eS+3/fzSbTKGy3wBAM2DMIJGSYqL0YNXn3xmzcrthyxWAwAIJ4QRNNp3sru77acWbLZYCQAgnBBG0GixMVG68MuH6H28/ZC27j9quSIAQDggjKBJfvT1Pm57/JNLGTsCADhvTQ4jy5Yt0zXXXKPMzEw5jqN58+adcf0lS5bIcZzTXsXFxedaMyyaNDgjaHrRpn2WKgEAhIsmh5Hy8nINHTpUs2bNatJ2BQUFKioqcl+pqalN/Wi0Ebd9rafbXru71GIlAIBwEN3UDSZNmqRJkyY1+YNSU1OVnJzc5O3Q9tw3sb9e/HC7JOkf+Xs048oLLFcEAAhlrTZmZNiwYcrIyNCVV16pjz766IzrVlVVye/3B73QdnijPRrSNUmStOPgMa3YdtByRQCAUNbiYSQjI0OzZ8/WG2+8oTfeeENZWVkaN26cPv300wa3yc3NVVJSkvvKyspq6TLRRI9eM9Bt3/DCCouVAABCnWPO43IIx3E0d+5cTZkypUnbjR07Vt26ddOf//znepdXVVWpqqrKnfb7/crKylJpaakSExPPtVw0sx73v+O2F90zVr06t7dYDQCgrfH7/UpKSjrr77eVS3tHjRqlLVu2NLjc5/MpMTEx6IW2Z8XM8W77/jfWW6wEABDKrISR/Px8ZWRknH1FtGnpSbFue/vBcouVAABCWZOvpjl69GjQUY3t27crPz9fHTt2VLdu3TRz5kzt2bNHf/rTnyRJTz/9tHr27KmBAweqsrJSL774ohYtWqT33nuv+XoBa+4Y21uzl27V/rIqFZdWBgUUAAAao8lHRlatWqXhw4dr+PDhkqQZM2Zo+PDheuSRRyRJRUVFKiwsdNevrq7WPffco8GDB2vs2LFau3at3n//fY0fP77e90douXNsb7d9Se5Ci5UAAELVeQ1gbS2NHQADO04dyLrj8ckWKwEAtCVtegArwsuqh3Lc9gdf7LdYCQAgFBFGcN5S2vvc9m/mb7JYCQAgFBFG0CxG9ewoSdqwx699ZZWWqwEAhBLCCJrF1/uffPDhqyt3WawEABBqCCNoFkO7Jrvt372/2V4hAICQQxhBs8ju3Sloetv+o5YqAQCEGsIIms3Lt17stp9cwNERAEDjEEbQbMb1Ozlu5J11RRYrAQCEEsIImtVlfU6ergmB++kBANoAwgia1fRxfdz21v08PA8AcHaEETSrUwey5jy11GIlAIBQQRhBs3IcJ2h6f1mVpUoAAKGCMIJm986Pvua2X/248AxrAgBAGEELGJiZ5Laf4hJfAMBZEEbQ4vYeqbBdAgCgDSOMoEX8Ysogt71o0z6LlQAA2jrCCFpEdq+ObvvX7260WAkAoK0jjKBF9ElNcNvHqussVgIAaOsII2gxr9w22m0v27zfYiUAgLaMMIIW0yetvdu+9eVPLFYCAGjLCCNoMakJsUrwRUuS6gJGdQGeVQMAOB1hBC1q1s0Xue19ZZUWKwEAtFWEEbSoMRd0dtsbi/wWKwEAtFWEEbS4hNjjp2q+9/Iqy5UAANoiwgha3OV9U9x2ZQ2X+QIAghFG0OJuu7yX2/7t/AKLlQAA2iLCCFrc0K7Jbvulj7bbKwQA0CYRRtDiojyOHr1mgDsd4BJfAMApCCNoFTdc3M1t5zy11GIlAIC2hjCCVhHnjXLb2w6UW6wEANDWEEZgxa5Dx2yXAABoIwgjaDXvzxjrti//7WKLlQAA2hLCCFpNn9T2QdMMZAUASIQRtLJnT3lWzZpdhy1WAgBoKwgjaFWTBqW77W89l2exEgBAW0EYQatyHCdo2hhO1QBApCOMoNXddUUft71ud6nFSgAAbQFhBK3unqsucNvXzvrIYiUAgLaAMIJW99VTNZ/t5egIAEQywgis+Otto9325P/7ocVKAAC2EUZgxWV9UmyXAABoIwgjsGZgZqLbLqussVgJAMAmwgiseem7F7vtn/3jc4uVAABsIozAmrTEWLf9xqe7Vcft4QEgIhFG0Gbc87d82yUAACwgjMCq135widuel7/XYiUAAFsII7BqdK9OQdPcHh4AIg9hBNZNy+7utnvOfFeVNXUWqwEAtDbCCKw7NYxI0rNLtlqqBABgA2EE1vVMaR80/cbq3ZYqAQDYQBiBdVEeRx/cd4U7fWFGgsVqAACtjTCCNiGrY7xG9ewoSXp/4z7tOnTMckUAgNZCGEGb8etvDnLb97+5zmIlAIDWRBhBm9G1Q7zb/mjLQYuVAABaE2EEbUZsTFTQ9JZ9ZZYqAQC0JsII2pRNv5jotnOeWmaxEgBAayGMoE356tGRqlpugAYA4Y4wgjbn3R9d7rYXbtxnsRIAQGsgjKDNGZCZ6LZ/+NdPLVYCAGgNhBG0eZuK/bZLAAC0oCaHkWXLlumaa65RZmamHMfRvHnzzrrNkiVLdNFFF8nn86lPnz56+eWXz6FURJKbRndz2796Z6PFSgAALa3JYaS8vFxDhw7VrFmzGrX+9u3bNXnyZF1xxRXKz8/X3Xffrdtuu03/+te/mlwsIsevppy8AdoHXxyQMcZiNQCAlhTd1A0mTZqkSZMmNXr92bNnq2fPnnryySclSRdeeKE+/PBD/e53v9OECROa+vGIEI7j6D/G9NLzy7ZJkpZu3q9x/VItVwUAaAktPmYkLy9POTk5QfMmTJigvLy8BrepqqqS3+8PeiHy3Dexv9t+asFmi5UAAFpSi4eR4uJipaWlBc1LS0uT3+9XRUVFvdvk5uYqKSnJfWVlZbV0mWiDojyORnbvIElat7tUdQFO1QBAOGqTV9PMnDlTpaWl7mvXrl22S4Ill/VJcdu3/c8nFisBALSUFg8j6enpKikpCZpXUlKixMRExcXF1buNz+dTYmJi0AuR6eZLTl5Vs7hgv8VKAAAtpcXDSHZ2thYuXBg0b8GCBcrOzm7pj0YYSE2IVa+UdrbLAAC0oCaHkaNHjyo/P1/5+fmSjl+6m5+fr8LCQknHT7FMmzbNXf+OO+7Qtm3bdN9992nTpk169tln9be//U0/+clPmqcHCHtv3Hmp2y49VmOxEgBAS2hyGFm1apWGDx+u4cOHS5JmzJih4cOH65FHHpEkFRUVucFEknr27Kl33nlHCxYs0NChQ/Xkk0/qxRdf5LJeNFqHdl4lxcVIkq5/frnlagAAzc0xIXA3Kb/fr6SkJJWWljJ+JEL1uP8dt73pFxNPe7ovAKDtaezvd5u8mgb4qqFZyW771Y8LG14RABByCCMICW/ccXLA82P/+7nFSgAAzY0wgpAQHcVXFQDCFf+FR8j4+bUD3XbhwWMWKwEANCfCCELGhIHpbnvNrsMWKwEANCfCCEJGWmKsuiQfv2vvHz/aYbcYAECzIYwgpHz30h6SpPxdR1RWyQ3QACAcEEYQUk49VfM/y3fYKwQA0GwIIwgpWR1PPlyxqLTSYiUAgOZCGEFIcRxHT/yfIZKkv64sVG1dwHJFAIDzRRhByBmYmeS2V+3kqhoACHWEEYScAZknn2/w4RcHLFYCAGgOhBGEpD6p7SVJzyzeYrkSAMD5IowgJH2tT4rb9nOJLwCENMIIQtKj1wxw2w+8ud5iJQCA80UYQUhyHMdtv72uSMYYi9UAAM4HYQQh68oBaW7bX1FrsRIAwPkgjCBkPXfzRW77f9fttVgJAOB8EEYQsqKjTn59H5q3wWIlAIDzQRgBAABWEUYQ0qaOzHLb5VWMGwGAUEQYQUh7/FuD3fa1sz6yWAkA4FwRRhDSTr3Ed8u+oxYrAQCcK8IIwkpdgPuNAECoIYwg5P3jrsvc9jvriyxWAgA4F4QRhLwhXZPd9o9eXaNj1QxkBYBQQhhBWLi878kH5729jqMjABBKCCMIC2Mv6Oy2P95+yGIlAICmIowgLHzvsp5u+++rd1usBADQVIQRhAWPx1FyfIw7XVFdZ7EaAEBTEEYQNv72H9lum6tqACB0EEYQNjKSYt32C8u2WqwEANAUhBGEjfa+aLe9ueQoz6oBgBBBGEHYcBxHz9w03J1eunm/xWoAAI1FGEFYmTQow23/4u3PLVYCAGgswgjCSpTHUWzM8a91UWml5WoAAI1BGEHYueXSHm578aZ99goBADQKYQRh56ZR3dz2f766xmIlAIDGIIwg7HTv1M5tH62qlTHGYjUAgLMhjCAsvfTdkW57+4Fyi5UAAM6GMIKw9PX+aYqLiZIkvfTRdsvVAADOhDCCsFVRc/z5NH9ZUWi5EgDAmRBGELZyrxvstveVcZkvALRVhBGErSnDurjtUb9aaLESAMCZEEYQtuK8UbZLAAA0AmEEYe3FaSevqtm2/6jFSgAADSGMIKzlDEhz219/cqnFSgAADSGMIOxdd9HJsSOVX15hAwBoOwgjCHv3T+rvtq97drnFSgAA9SGMIOx1bu9z258X+S1WAgCoD2EEYc9xHGUmxbrT+8uqLFYDAPgqwggiwkf3f91tj3tiscVKAABfRRhBRHAcx22XV9fxJF8AaEMII4gYz39nhNteuHGfxUoAAKcijCBiXNq7k9v+x9q9FisBAJyKMIKIkRAbo1sv6yHpeBjhVA0AtA2EEUSU2y/v5bYfn7/JYiUAgBMII4gomclxbvv5pdu0+/Axi9UAACTCCCLQqXdk/Y8/r7ZYCQBAIowgAv3glFM1n+3ljqwAYNs5hZFZs2apR48eio2N1ejRo/Xxxx83uO7LL78sx3GCXrGxsQ2uD7Q0j8fRY98Y6E5/WnjYYjUAgCaHkddee00zZszQo48+qk8//VRDhw7VhAkTtG9fw/dtSExMVFFRkfvauXPneRUNnK+v9U1x29c9u5wrawDAoiaHkaeeekq33367br31Vg0YMECzZ89WfHy8XnrppQa3cRxH6enp7istLe28igbOV+/O7XX5KYHkjU/3WKwGACJbk8JIdXW1Vq9erZycnJNv4PEoJydHeXl5DW539OhRde/eXVlZWbr22mv12WefnfFzqqqq5Pf7g15Ac5v97ZN3ZP2v19darAQAIluTwsiBAwdUV1d32pGNtLQ0FRcX17tNv3799NJLL+mtt97SX/7yFwUCAV166aXavXt3g5+Tm5urpKQk95WVldWUMoFGaeeLVrTn5DNrPtlxyGI1ABC5WvxqmuzsbE2bNk3Dhg3T2LFj9eabb6pz5856/vnnG9xm5syZKi0tdV+7du1q6TIRoZbdd4Xbvn52HmNHAMCCJoWRlJQURUVFqaSkJGh+SUmJ0tPTG/UeMTExGj58uLZs2dLgOj6fT4mJiUEvoCVkJscpOT7Gnf7nhvqP8AEAWk6TwojX69WIESO0cOFCd14gENDChQuVnZ3dqPeoq6vT+vXrlZGR0bRKgRay4Cdj3fYP//qpxUoAIDJFN3WDGTNm6JZbbtHIkSM1atQoPf300yovL9ett94qSZo2bZq6dOmi3NxcSdLPf/5zXXLJJerTp4+OHDmiJ554Qjt37tRtt93WvD0BzlHnBF/QdGlFjZLiYhpYGwDQ3JocRqZOnar9+/frkUceUXFxsYYNG6b58+e7g1oLCwvl8Zw84HL48GHdfvvtKi4uVocOHTRixAgtX75cAwYMaL5eAOfp+e+McG8NP/zn72lb7mTLFQFA5HBMCIzY8/v9SkpKUmlpKeNH0CKMMeo58113eu0jVykpnqMjAHA+Gvv7zbNpAB2/Md/f7zg57unb//9Ki9UAQGQhjABfGtmjo9tev6dUh8qrLVYDAJGDMAKcYum949z2Jb9e2PCKAIBmQxgBTtG9Uzt9Y2imJKm6LqC9RyosVwQA4Y8wAnzF01OHue1LH1+kqto6e8UAQAQgjABf4fE4ivdGudNPzC+wWA0AhD/CCFCPjx88+WTqFz/crrpAm78CHgBCFmEEqEd7X7Rmf3uEOz3nk0KL1QBAeCOMAA2YMDDNbT84d4M2l5RZrAYAwhdhBGiA4zj6xZRB7vRVv1tmsRoACF+EEeAMvj26W9D00apaS5UAQPgijABn4DiOfn/DMHd60KP/Ugg8zgkAQgphBDiLa4ZkBk0/v2ybpUoAIDwRRoCz8Hgcbf7lJHf68X9u0oGjVRYrAoDwQhgBGsEb7dE9V17gTo/85fucrgGAZkIYARrpjnG9g6anzPrIUiUAEF4II0AjxUR5tOPxye702t2lOsjpGgA4b4QRoIn+dfcYtz3il++rsoYH6QHA+SCMAE3ULz0haLr/w/MV4Nk1AHDOCCPAOVjwkzFB0z/5W76dQgAgDBBGgHPQNy1Bc394qTv9Vv5erdh20GJFABC6CCPAORrerYN6dIp3p294YYU+2XHIYkUAEJoII8B5WHTPuKDp62fncf8RAGgiwghwHjweR1t+NSloXs+Z71qqBgBCE2EEOE/RUR79711fC5r3zKIvLFUDAKGHMAI0g8Fdk7TwnrHu9P/33mYt3rRPdVzyCwBnRRgBmknvzu31u6lD3elbX/5ET/yrwGJFABAaCCNAM/rm8K763mU93enZS7dqxmv59goCgBBAGAGa2SPXDNDj1w12p99cs0f/WLvXYkUA0LYRRoAWcMOobhqalexO/+jVNdwUDQAaQBgBWshb0y9TtMdxp294YYX6PfRPixUBQNtEGAFa0OZfBt+DpKo2oJv+ewVX2QDAKQgjQAvyeBx9/vMJQfOWbz2o3g+8q4rqOktVAUDbQhgBWli8N1o7Hp+s+yf1D5p/4SPzVVRaYakqAGg7CCNAK7ljbG/9/oZhQfOycxfp2mc+VFlljZ2iAKANIIwArejaYV30wX1XBM1bu7tUg3/2nqWKAMA+wgjQyrI6xp82jkSSetz/jg4erbJQEQDYRRgBLIj3RuuLX03S1YPTg+aP+OX7+sXbn1uqCgDscIwxbf4aQ7/fr6SkJJWWlioxMdF2OUCz2rKvTDlPLTtt/icP5qhzgs9CRQDQPBr7+82REcCyPqkJ2vrrq5VzYWrQ/It/9b563P+OKmu4BBhAeCOMAG1AlMfRf08bqQeu7n/asv4Pz9d7nxVbqAoAWgenaYA2pqYuoL4P1n/b+E8fvlId23lbuSIAODeN/f0mjABtUCBglL/7iK57dnm9y9f97Colxsa0clUA0DSMGQFCmMfj6KJuHbTygfH1Lh/ys/f0jWc+1KZifytXBgDNjyMjQAgwxuiuV9fonXVF9S5/6bsjNbpnJ7XzRbdyZQDQME7TAGGoujagCx6qfzyJJN1+eU/dO6G/vNEc9ARgH2EECGP+yhr9+NU1Wlywv8F1Zlx5ge66oo88HqcVKwOAkwgjQAQIBIz+J2+HHvvfM9+1dfa3R2jioPQzrgMAzY0wAkSYnQfLNfaJJWdc56bR3XR3Tl+lJsS2TlEAIhphBIhQG4v8evK9zXp/Y8kZ1xuYmag/3nqx2nmjGfgKoEUQRgBIklZuO6ipL6xo1Lq//uZg3TgqS47DOBMA548wAiBIZU2dPt15WPe9sU67D1ecdf0bR3XTf111gTq28xJOAJwTwgiAMyoqrVDOk0tVXt34B/H9nxFdNWFgusb166yYKC4fBnBmhBEAjVZZU6c1hUd043837nTOV+XN/Lpqao26dojjUmIALsIIgPOy90iF/rJip/6ct1NlVbVN3t4X7dG9E/opMTZG3xrRVbWBgHzRUS1QKYC2ijACoFkFAkbbDpTr+tnLdfhYzXm9V0yUo2uGZOrKAWlyHEfj+nVWbAxBBQg3hBEArWJfWaX+tHyn/rxip0orzi+kfNWlvTtpWnYPeaMdjezRUfExUYpmrAoQMggjAKw6VF6td9btVbw3WosK9jX4kL/mMmFgmv59ZJaiozwamJkob7RH3iiPfNEeOY4jYwxXBQGtjDACoE07cqxaK7cfUmlFjd5dX6TCg8e07UB5q9fRzhulK/qnavLgDG07UK70xFhd3KOjYmM8Sk2MVSBg5DiS4zgKBAwDdIEmIIwACHlHq2q1bvcROXJU4q/Ub+ZvUlFppdWa2vuiday6VqN7dlJiXLT+9VmJUtr7VFVbJ2Okfx+ZpfkbinSkokbXXdRFEwamq7ImoOgoR71S2qlDO68qa+qU0s6nOmMU7XE4YoOw1aJhZNasWXriiSdUXFysoUOH6g9/+INGjRrV4Pqvv/66Hn74Ye3YsUN9+/bVb37zG1199dWN/jzCCIDGqKiu0wdf7Nem4jJ9suOQ/BU1Wru71HZZzebiHh3Uu3N7bT9QrpXbD0mSLuqWrPKqOrXzRakuYJSRFKesjnHKSIqTJKUnxapfeoI27CnVkK7J2n34mBJiY9S1Q5zaeaO1t7RCvVLaKWCkKI+j8qpaxcZEKYojQGgGLRZGXnvtNU2bNk2zZ8/W6NGj9fTTT+v1119XQUGBUlNTT1t/+fLlGjNmjHJzc/Vv//ZveuWVV/Sb3/xGn376qQYNGtSsnQGAcxEIGNUEAqoLGNUGjEqP1ag2YLT9wFEtLdivRQX7dFG3DiooLtPonh11tKpOVbV1WrHtkA4crbJdfovwOJLHcVQbOPkT4Y32aHCXJK3eeVgDMxP12V7/adv17txOlTUBSVJFTZ2iPI4mD87QX1bsVG3A6IK09hqYmaS5a/boxlHddORYteoCRr6YKG3Zd1QeR4qNidKVA9KUEButeG+U1u/2a++RCg3JStKKbYc0snsH9UtP0KodhzSka7LmbyiWx+Ooc3ufRnTvoHhvlOK8USrxV6pjO68ykuJ0qLxajiPFRHlUFwgoLiZaiXHR8kZ5VFpRo5T2PtUGjHwxHpUeq1Gn9l7Fe48/s+nEz6TjOKoLGJ2IaUaSI5126s4YI2NOnx+JWiyMjB49WhdffLGeeeYZSVIgEFBWVpb+8z//U/fff/9p60+dOlXl5eV6++233XmXXHKJhg0bptmzZzdrZwCgLQkEjPuDZSQdLK/S4fIa5e86rLLKWnVO8Kmmzqig2K/NJUclSUs37w96jxM/+u28URrcNUnlVXVavyd8jva0dY4jNeZXMjk+RpU1dYqNidKRJlz67o32qLo2oC7JcUqIjVZZZa3SEn0qrajR1v3HxzAV+yvVOcGnzOQ4xUZ75DiSNzpKpceqtb+sSnu/PHXZrWO8EuOiVXSkUqUVNeqblqAuyXHyRjuqqK7TniMVSoqLUZfk40fN6szx72hdwGhAZqLuHNe72e+s3Njf7yY9qrO6ulqrV6/WzJkz3Xkej0c5OTnKy8urd5u8vDzNmDEjaN6ECRM0b968Bj+nqqpKVVUn/2/D7z89fQNAW/fV/zNOTYhVasLx0yatrbyqVvHeKB0+VqN4b5SMkbYdOKqq2oA+3+tXdW1Amcmx8lfUqqyqVt06xkuS/rpyp0Z276Ct+8vVM6Wd4r1R2n6gXJ0TfNp58JjmrtmjfmkJKigpkzfao8v7pMgb7dH2A+WKifIoMzlW8d5ord9TqijHkeNIe45UqKyyVqN7dlRiXIy27jt62uDlBF+0hnVL1sGj1fq86PTfgC7Jcdpz5OzPWDpfjf3f9RMB5MRRocaqrj2+/ql9ObVd7D8eNPaXVWl/2ZmPwhUeOhY0vbHIr431/O0+0eHT5s3/rFgp7X26aXS3xhffjJoURg4cOKC6ujqlpaUFzU9LS9OmTZvq3aa4uLje9YuLixv8nNzcXD322GNNKQ0AcAbtfMf/c9+xndedNzAzSZJ0UbcODW535YC0BpdJ0u+mDjv/4ppZQ5dxn7gySpJqA0Yex1FlTZ3ivVGq/fIIQcAc/+e+sirFxkTJG+WRx5G2HyhXcrxX7XxRKq+q1T5/lTq296q2zsgXffxUT/vYaHkcR4fKq+WvqFFFTZ3ivryZnzfao3W7S5UcH6M9hyskR+qaHKfUxFit2nFIYy7orMqagA4crVKU46hDO6+qautUXlWrwkPHlNUhXhU1dUqKi1F0lEc1tQFV1wW0aNM+ZSTFSjr+WIfOCT7tOXz8CEhSXIzivMdPdTmOVFZZq6ragOK9UYpyHHk8jqI9jopKK3XkWLVuuDir1fbRVzUpjLSWmTNnBh1N8fv9ysqy90cCAISOhq5OOvVIVUzU8faJkBYT5ejUmwAnxMYEbdupvS9ouk9q049ujet3+rhKSZowML3J73XCjaPsHMlobk0KIykpKYqKilJJSUnQ/JKSEqWn1//HTE9Pb9L6kuTz+eTz+RpcDgAAwkeTRqp4vV6NGDFCCxcudOcFAgEtXLhQ2dnZ9W6TnZ0dtL4kLViwoMH1AQBAZGnyaZoZM2bolltu0ciRIzVq1Cg9/fTTKi8v16233ipJmjZtmrp06aLc3FxJ0o9//GONHTtWTz75pCZPnqw5c+Zo1apVeuGFF5q3JwAAICQ1OYxMnTpV+/fv1yOPPKLi4mINGzZM8+fPdwepFhYWyuM5ecDl0ksv1SuvvKKHHnpIDzzwgPr27at58+Y1+h4jAAAgvHE7eAAA0CIa+/vNs7gBAIBVhBEAAGAVYQQAAFhFGAEAAFYRRgAAgFWEEQAAYBVhBAAAWEUYAQAAVrXJp/Z+1Yn7svn9fsuVAACAxjrxu322+6uGRBgpKyuTJGVlZVmuBAAANFVZWZmSkpIaXB4St4MPBALau3evEhIS5DhOs72v3+9XVlaWdu3aFTG3mafP9Dlc0Wf6HK5Cuc/GGJWVlSkzMzPouXVfFRJHRjwej7p27dpi75+YmBhyO/h80efIQJ8jA32ODKHa5zMdETmBAawAAMAqwggAALAqosOIz+fTo48+Kp/PZ7uUVkOfIwN9jgz0OTJEQp9DYgArAAAIXxF9ZAQAANhHGAEAAFYRRgAAgFWEEQAAYFVEh5FZs2apR48eio2N1ejRo/Xxxx/bLumc/OxnP5PjOEGv/v37u8srKys1ffp0derUSe3bt9e3vvUtlZSUBL1HYWGhJk+erPj4eKWmpuree+9VbW1ta3elQcuWLdM111yjzMxMOY6jefPmBS03xuiRRx5RRkaG4uLilJOToy+++CJonUOHDunmm29WYmKikpOT9f3vf19Hjx4NWmfdunW6/PLLFRsbq6ysLP32t79t6a416Gx9/u53v3vafp84cWLQOqHW59zcXF188cVKSEhQamqqpkyZooKCgqB1muv7vGTJEl100UXy+Xzq06ePXn755ZbuXr0a0+dx48adtq/vuOOOoHVCqc/PPfechgwZ4t7EKzs7W//85z/d5eG2j6Wz9znc9nGTmQg1Z84c4/V6zUsvvWQ+++wzc/vtt5vk5GRTUlJiu7Qme/TRR83AgQNNUVGR+9q/f7+7/I477jBZWVlm4cKFZtWqVeaSSy4xl156qbu8trbWDBo0yOTk5Jg1a9aYd99916SkpJiZM2fa6E693n33XfPggw+aN99800gyc+fODVr++OOPm6SkJDNv3jyzdu1a841vfMP07NnTVFRUuOtMnDjRDB061KxYscJ88MEHpk+fPubGG290l5eWlpq0tDRz8803mw0bNphXX33VxMXFmeeff761uhnkbH2+5ZZbzMSJE4P2+6FDh4LWCbU+T5gwwfzxj380GzZsMPn5+ebqq6823bp1M0ePHnXXaY7v87Zt20x8fLyZMWOG+fzzz80f/vAHExUVZebPn9+q/TWmcX0eO3asuf3224P2dWlpqbs81Pr8j3/8w7zzzjtm8+bNpqCgwDzwwAMmJibGbNiwwRgTfvvYmLP3Odz2cVNFbBgZNWqUmT59ujtdV1dnMjMzTW5ursWqzs2jjz5qhg4dWu+yI0eOmJiYGPP666+78zZu3Ggkmby8PGPM8R89j8djiouL3XWee+45k5iYaKqqqlq09nPx1R/mQCBg0tPTzRNPPOHOO3LkiPH5fObVV181xhjz+eefG0nmk08+cdf55z//aRzHMXv27DHGGPPss8+aDh06BPX5pz/9qenXr18L9+jsGgoj1157bYPbhHqfjTFm3759RpJZunSpMab5vs/33XefGThwYNBnTZ061UyYMKGlu3RWX+2zMcd/qH784x83uE2o99kYYzp06GBefPHFiNjHJ5zoszGRsY/PJCJP01RXV2v16tXKyclx53k8HuXk5CgvL89iZefuiy++UGZmpnr16qWbb75ZhYWFkqTVq1erpqYmqK/9+/dXt27d3L7m5eVp8ODBSktLc9eZMGGC/H6/Pvvss9btyDnYvn27iouLg/qYlJSk0aNHB/UxOTlZI0eOdNfJycmRx+PRypUr3XXGjBkjr9frrjNhwgQVFBTo8OHDrdSbplmyZIlSU1PVr18/3XnnnTp48KC7LBz6XFpaKknq2LGjpOb7Pufl5QW9x4l12sK//1/t8wl//etflZKSokGDBmnmzJk6duyYuyyU+1xXV6c5c+aovLxc2dnZEbGPv9rnE8J1HzdGSDwor7kdOHBAdXV1QTtVktLS0rRp0yZLVZ270aNH6+WXX1a/fv1UVFSkxx57TJdffrk2bNig4uJieb1eJScnB22Tlpam4uJiSVJxcXG9f4sTy9q6EzXW14dT+5iamhq0PDo6Wh07dgxap2fPnqe9x4llHTp0aJH6z9XEiRN13XXXqWfPntq6daseeOABTZo0SXl5eYqKigr5PgcCAd1999267LLLNGjQILem5vg+N7SO3+9XRUWF4uLiWqJLZ1VfnyXppptuUvfu3ZWZmal169bppz/9qQoKCvTmm29KCs0+r1+/XtnZ2aqsrFT79u01d+5cDRgwQPn5+WG7jxvqsxSe+7gpIjKMhJtJkya57SFDhmj06NHq3r27/va3v7XpLx/Ozw033OC2Bw8erCFDhqh3795asmSJxo8fb7Gy5jF9+nRt2LBBH374oe1SWk1Dff7BD37gtgcPHqyMjAyNHz9eW7duVe/evVu7zGbRr18/5efnq7S0VH//+991yy23aOnSpbbLalEN9XnAgAFhuY+bIiJP06SkpCgqKuq00dklJSVKT0+3VFXzSU5O1gUXXKAtW7YoPT1d1dXVOnLkSNA6p/Y1PT293r/FiWVt3Ykaz7Q/09PTtW/fvqDltbW1OnToUNj8HXr16qWUlBRt2bJFUmj3+a677tLbb7+txYsXq2vXru785vo+N7ROYmKitQDfUJ/rM3r0aEkK2teh1mev16s+ffpoxIgRys3N1dChQ/X73/8+rPdxQ32uTzjs46aIyDDi9Xo1YsQILVy40J0XCAS0cOHCoPN3oero0aPaunWrMjIyNGLECMXExAT1taCgQIWFhW5fs7OztX79+qAfrgULFigxMdE9hNiW9ezZU+np6UF99Pv9WrlyZVAfjxw5otWrV7vrLFq0SIFAwP2XPjs7W8uWLVNNTY27zoIFC9SvX782d4qmPrt379bBgweVkZEhKTT7bIzRXXfdpblz52rRokWnnUJqru9zdnZ20HucWMfGv/9n63N98vPzJSloX4dSn+sTCARUVVUVlvu4ISf6XJ9w3MdnZHsErS1z5swxPp/PvPzyy+bzzz83P/jBD0xycnLQSOVQcc8995glS5aY7du3m48++sjk5OSYlJQUs2/fPmPM8cvkunXrZhYtWmRWrVplsrOzTXZ2trv9iUvGrrrqKpOfn2/mz59vOnfu3KYu7S0rKzNr1qwxa9asMZLMU089ZdasWWN27txpjDl+aW9ycrJ56623zLp168y1115b76W9w4cPNytXrjQffvih6du3b9BlrkeOHDFpaWnmO9/5jtmwYYOZM2eOiY+Pt3aZ65n6XFZWZv7rv/7L5OXlme3bt5v333/fXHTRRaZv376msrLSfY9Q6/Odd95pkpKSzJIlS4IucTx27Ji7TnN8n09cAnnvvfeajRs3mlmzZlm7BPJsfd6yZYv5+c9/blatWmW2b99u3nrrLdOrVy8zZsyYkO3z/fffb5YuXWq2b99u1q1bZ+6//37jOI557733jDHht4+NOXOfw3EfN1XEhhFjjPnDH/5gunXrZrxerxk1apRZsWKF7ZLOydSpU01GRobxer2mS5cuZurUqWbLli3u8oqKCvPDH/7QdOjQwcTHx5tvfvObpqioKOg9duzYYSZNmmTi4uJMSkqKueeee0xNTU1rd6VBixcvNpJOe91yyy3GmOOX9z788MMmLS3N+Hw+M378eFNQUBD0HgcPHjQ33nijad++vUlMTDS33nqrKSsrC1pn7dq15mtf+5rx+XymS5cu5vHHH2+tLp7mTH0+duyYueqqq0znzp1NTEyM6d69u7n99ttPC9Oh1uf6+ivJ/PGPf3TXaa7v8+LFi82wYcOM1+s1vXr1CvqM1nS2PhcWFpoxY8aYjh07Gp/PZ/r06WPuvffeoHtQGBNaff7e975nunfvbrxer+ncubMZP368G0SMCb99bMyZ+xyO+7ipHGOMab3jMAAAAMEicswIAABoOwgjAADAKsIIAACwijACAACsIowAAACrCCMAAMAqwggAALCKMAIAAKwijAAAAKsIIwAAwCrCCAAAsIowAgAArPp/fP14qWN7ZbEAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# stage 1: train to 16 bits\n",
    "randkey, trainkey = jax.random.split(randkey)\n",
    "model = train_stage(trainkey, model, optimizer, batch_size,\n",
    "                    iterations=16384, min_val=0, max_val=1 << 16)\n",
    "test_fizzbuzz(model, [i for i in range(0, 25)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b109fb31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1961c00900d489fb4fc3e64dab44aed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16384 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100  fizz: 0.0000, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0001\n",
      "101  fizz: 0.0006, fizzmod: 2, buzz: 0.0000, buzzmod: 1, neither: 0.9998\n",
      "102  fizz: 1.0000, fizzmod: 0, buzz: 0.0008, buzzmod: 2, neither: 0.0000\n",
      "103  fizz: 0.0000, fizzmod: 1, buzz: 0.0001, buzzmod: 3, neither: 1.0000\n",
      "104  fizz: 0.0002, fizzmod: 2, buzz: 0.0001, buzzmod: 4, neither: 0.9990\n",
      "105  fizz: 0.9999, fizzmod: 0, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "106  fizz: 0.0003, fizzmod: 1, buzz: 0.0001, buzzmod: 1, neither: 0.9992\n",
      "107  fizz: 0.0001, fizzmod: 2, buzz: 0.0000, buzzmod: 2, neither: 0.9998\n",
      "108  fizz: 1.0000, fizzmod: 0, buzz: 0.0006, buzzmod: 3, neither: 0.0000\n",
      "109  fizz: 0.0001, fizzmod: 1, buzz: 0.0001, buzzmod: 4, neither: 1.0000\n",
      "110  fizz: 0.0000, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0001\n",
      "111  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 1, neither: 0.0000\n",
      "112  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 2, neither: 1.0000\n",
      "113  fizz: 0.0013, fizzmod: 2, buzz: 0.0000, buzzmod: 3, neither: 0.9996\n",
      "114  fizz: 1.0000, fizzmod: 0, buzz: 0.0009, buzzmod: 4, neither: 0.0000\n",
      "115  fizz: 0.0001, fizzmod: 1, buzz: 0.9998, buzzmod: 0, neither: 0.0008\n",
      "116  fizz: 0.0004, fizzmod: 2, buzz: 0.0001, buzzmod: 1, neither: 0.9997\n",
      "117  fizz: 0.9954, fizzmod: 0, buzz: 0.0000, buzzmod: 2, neither: 0.0259\n",
      "118  fizz: 0.0000, fizzmod: 1, buzz: 0.0001, buzzmod: 3, neither: 0.9999\n",
      "119  fizz: 0.0004, fizzmod: 2, buzz: 0.0001, buzzmod: 4, neither: 0.9991\n",
      "120  fizz: 0.9999, fizzmod: 0, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "121  fizz: 0.0001, fizzmod: 1, buzz: 0.0000, buzzmod: 1, neither: 0.9997\n",
      "122  fizz: 0.0001, fizzmod: 2, buzz: 0.0001, buzzmod: 2, neither: 0.9990\n",
      "123  fizz: 0.9999, fizzmod: 0, buzz: 0.0000, buzzmod: 3, neither: 0.0003\n",
      "124  fizz: 0.0000, fizzmod: 1, buzz: 0.0001, buzzmod: 4, neither: 1.0000\n",
      "1000000  fizz: 0.0000, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0001\n",
      "1000001  fizz: 0.0003, fizzmod: 2, buzz: 0.0000, buzzmod: 1, neither: 0.9999\n",
      "1000002  fizz: 1.0000, fizzmod: 0, buzz: 0.0003, buzzmod: 2, neither: 0.0000\n",
      "1000003  fizz: 0.0000, fizzmod: 1, buzz: 0.0001, buzzmod: 3, neither: 0.9999\n",
      "1000004  fizz: 0.0003, fizzmod: 2, buzz: 0.0006, buzzmod: 4, neither: 0.9988\n",
      "1000005  fizz: 0.9983, fizzmod: 0, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "1000006  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 1, neither: 0.9998\n",
      "1000007  fizz: 0.0001, fizzmod: 2, buzz: 0.0003, buzzmod: 2, neither: 0.9984\n",
      "1000008  fizz: 1.0000, fizzmod: 0, buzz: 0.0002, buzzmod: 3, neither: 0.0000\n",
      "1000009  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 4, neither: 1.0000\n",
      "1000010  fizz: 0.0004, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "1000011  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 1, neither: 0.0000\n",
      "1000012  fizz: 0.0000, fizzmod: 1, buzz: 0.0009, buzzmod: 2, neither: 0.9999\n",
      "1000013  fizz: 0.0001, fizzmod: 2, buzz: 0.0001, buzzmod: 3, neither: 0.9999\n",
      "1000014  fizz: 1.0000, fizzmod: 0, buzz: 0.0006, buzzmod: 4, neither: 0.0000\n",
      "1000015  fizz: 0.0000, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "1000016  fizz: 0.0002, fizzmod: 2, buzz: 0.0001, buzzmod: 1, neither: 0.9999\n",
      "1000017  fizz: 0.9994, fizzmod: 0, buzz: 0.0002, buzzmod: 2, neither: 0.0005\n",
      "1000018  fizz: 0.0001, fizzmod: 1, buzz: 0.0003, buzzmod: 3, neither: 0.9998\n",
      "1000019  fizz: 0.0001, fizzmod: 2, buzz: 0.0003, buzzmod: 4, neither: 0.9993\n",
      "1000020  fizz: 0.9996, fizzmod: 0, buzz: 0.9999, buzzmod: 0, neither: 0.0000\n",
      "1000021  fizz: 0.0027, fizzmod: 1, buzz: 0.0001, buzzmod: 1, neither: 0.9936\n",
      "1000022  fizz: 0.0001, fizzmod: 2, buzz: 0.0002, buzzmod: 2, neither: 0.9994\n",
      "1000023  fizz: 0.9996, fizzmod: 0, buzz: 0.0002, buzzmod: 3, neither: 0.0004\n",
      "1000024  fizz: 0.0001, fizzmod: 1, buzz: 0.0001, buzzmod: 4, neither: 0.9999\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhYAAAGfCAYAAAADEJteAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAsyUlEQVR4nO3de3hU1b3/8c+emWQSIJlwCxAJAnITEERAGvFyWlCLeKm//jw+HFpRe9paY4+XnlZiq9VTaWjt8adVS70VPbWCl6dgjwoWQUAqyEVQUMEbSKpcVEwm4TJJZtbvj0l2GAiYCTOzyM779Tz7mZk9a2Z/Zw8wH9Zae2/HGGMEAACQAj7bBQAAAO8gWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQhWAAAgJQJJNM4Go3q9ttv1xNPPKGdO3eqqKhIV155pX7xi1/IcZwWvUcsFtOnn36qvLy8Fr8GAADYZYxRdXW1ioqK5PMdpV/CJGHGjBmma9eu5vnnnzdbt241zzzzjOnUqZO59957W/weFRUVRhILCwsLCwtLG1wqKiqO+jufVI/Fa6+9pksuuUSTJ0+WJPXt21dz5szR6tWrW/weeXl5kqSKigrl5+cns3kAAGBJOBxWcXGx+zt+JEkFizPOOEMPPfSQ3nvvPQ0aNEhvvvmmVqxYobvvvvuIr4lEIopEIu7j6upqSVJ+fj7BAgCANuarpjEkFSymT5+ucDisIUOGyO/3KxqNasaMGZo6deoRX1NeXq477rgjmc0AAIA2KqmjQp5++mn95S9/0ZNPPqk33nhDjz/+uH73u9/p8ccfP+JrysrKVFVV5S4VFRXHXDQAADg+OcYY09LGxcXFmj59ukpLS911d955p5544glt3ry5Re8RDocVCoVUVVXFUAgAAG1ES3+/k+qx2Ldv32GHmPj9fsVisdZVCQAAPCWpORYXXXSRZsyYoT59+mjYsGFav3697r77bl199dXpqg8AALQhSQ2FVFdX69Zbb9W8efO0e/duFRUVacqUKbrtttuUnZ3dovdgKAQAgLanpb/fSQWLVCBYAADQ9qRljgUAAMDRECwAAEDKECwAAEDKECwAAEDKECwAAEDKJHUei+PZ3X/fovCBel1zzknqGcqxXQ4AAO2SZ3os5qyp0GOvbdOevbW2SwEAoN3yTLBovIirUUZPywEAAA7inWDRkCwye7ovAABwMO8EC7fPAgAA2OKdYEGPBQAA1nknWDTcMscCAAB7vBMsGros6LEAAMAezwSLRuQKAADs8UywaJpjQbQAAMAW7wULu2UAANCueSdYiDkWAADY5p1g4Z7GgmQBAIAt3gkWDbf0WAAAYI93gkXj4aaW6wAAoD3zTrBouKXHAgAAezwTLMThpgAAWOeZYNF0Sm8AAGCLd4IFp/QGAMA67wSLhlsuQgYAgD3eCRaMhQAAYJ13goU43BQAANu8Eyzco0Ls1gEAQHvmmWDRiDkWAADY45lgwVEhAADY551g0XBLrgAAwB7vBAvOvAkAgHXeCxZ2ywAAoF1LKlj07dtXjuMctpSWlqarvhZzRLIAAMC2QDKN16xZo2g06j7etGmTzj33XF122WUpLyxZTT0WJAsAAGxJKlh079494fHMmTN10kkn6ZxzzjniayKRiCKRiPs4HA4nWWLLcNl0AADsa/Uci9raWj3xxBO6+uqr3UM9m1NeXq5QKOQuxcXFrd3k0XG4KQAA1rU6WMyfP1+VlZW68sorj9qurKxMVVVV7lJRUdHaTR4Vh5sCAGBfUkMhB3v00Uc1adIkFRUVHbVdMBhUMBhs7WZajMNNAQCwr1XB4uOPP9bLL7+sv/71r6mup9XosQAAwL5WDYXMnj1bhYWFmjx5cqrraTVO6Q0AgH1JB4tYLKbZs2dr2rRpCgRaPZKSck3TR0kWAADYknSwePnll7V9+3ZdffXV6agHAAC0YUl3OZx33nnH5QTJpsmbdusAAKA98861QhoGQ8gVAADY45lgIXosAACwzjPBoulwU5IFAAC2eCdY0GMBAIB13gkWzLEAAMA67wQLTukNAIB1ngsWAADAHu8EC3FKbwAAbPNOsGgcCmGWBQAA1ngmWDSixwIAAHs8Eyy4uikAAPZ5J1g03JIrAACwxzvBgsNNAQCwzjvBouGWWAEAgD3eCRZNh4UAAABLvBMsGm453BQAAHu8Eyy4CBkAANZ5JliIi5ABAGCdZ4IFPRYAANjnnWDRcMscCwAA7PFOsKDHAgAA67wTLJhjAQCAdd4JFu5YCNECAABbPBcsiBUAANjjnWAhrm4KAIBtngkW4iJkAABY55lgwUXIAACwzzvBwmEoBAAA27wTLBpuyRUAANjjnWDBHAsAAKzzTrCwXQAAAPBQsGCOBQAA1nknWDTcchEyAADsSTpYfPLJJ/rOd76jrl27Kjc3V6eccorWrl2bjtqSw0XIAACwLpBM4y+//FLjx4/X17/+dS1YsEDdu3fX+++/r86dO6ervhbjImQAANiXVLD4zW9+o+LiYs2ePdtd169fv6O+JhKJKBKJuI/D4XCSJbYMl00HAMC+pIZC/va3v2nMmDG67LLLVFhYqFGjRunhhx8+6mvKy8sVCoXcpbi4+JgKPhLmWAAAYF9SweKjjz7SrFmzNHDgQL300kv60Y9+pP/4j//Q448/fsTXlJWVqaqqyl0qKiqOuejm0GMBAIB9SQ2FxGIxjRkzRr/+9a8lSaNGjdKmTZv0xz/+UdOmTWv2NcFgUMFg8Ngr/QoOZ7IAAMC6pHosevXqpaFDhyasO/nkk7V9+/aUFtUanHkTAAD7kgoW48eP15YtWxLWvffeezrxxBNTWlRrMBQCAIB9SQWLG2+8UatWrdKvf/1rffDBB3ryySf10EMPqbS0NF31JYHDTQEAsC2pYDF27FjNmzdPc+bM0fDhw/WrX/1K99xzj6ZOnZqu+lqMHgsAAOxLavKmJF144YW68MIL01HLMeFwUwAA7PPOtULosQAAwDrvBAvmWAAAYJ13goU7FkK0AADAFu8Ei4ZbYgUAAPZ4J1g0dFnQYQEAgD2eCRaNOCoEAAB7PBMsOCoEAAD7vBMsOCoEAADrvBMs6LEAAMA67wSLhlvmWAAAYI93ggXHmwIAYJ2HggVzLAAAsM07waLh1jDJAgAAazwTLMTkTQAArPNMsOBwUwAA7PNOsKDHAgAA67wTLBpuOdwUAAB7vBMs6LEAAMA67wQLt88CAADY4p1g4fZY0GUBAIAtngkWAADAPs8EC87oDQCAfZ4JFo1jIYyEAABgj2eCBYebAgBgn3eCBYebAgBgnXeCBaf0BgDAOu8EC3osAACwzjvBwr1HsgAAwBbvBAt6LAAAsM5DwYLDTQEAsM0zwaIRh5sCAGCPZ4IFQyEAANiXVLC4/fbb5ThOwjJkyJB01ZYUDjcFAMC+QLIvGDZsmF5++eWmNwgk/RZpQY8FAAD2JZ0KAoGAevbsmY5ajgmn9AYAwL6k51i8//77KioqUv/+/TV16lRt3779qO0jkYjC4XDCkg4OlzcFAMC6pILFuHHj9Nhjj2nhwoWaNWuWtm7dqrPOOkvV1dVHfE15eblCoZC7FBcXH3PRzWGOBQAA9iUVLCZNmqTLLrtMI0aM0Pnnn68XX3xRlZWVevrpp4/4mrKyMlVVVblLRUXFMRfdnKY5FkQLAABsOaaZlwUFBRo0aJA++OCDI7YJBoMKBoPHspmkECsAALDnmM5jUVNTow8//FC9evVKVT2txpk3AQCwL6lg8Z//+Z9atmyZtm3bptdee02XXnqp/H6/pkyZkq76Woy5mwAA2JfUUMg///lPTZkyRV988YW6d++uM888U6tWrVL37t3TVV+LMccCAAD7kgoWc+fOTVcdx4weCwAA7PPMtUJ8vsY5FkQLAABs8UywYPImAAD2eSZYNHRYKBojWQAAYItngoW/oceCXAEAgD2eCRY+hzkWAADY5plg0Xi4aZRgAQCANZ4JFn4fQyEAANjmmWDBUAgAAPZ5Jlg4HBUCAIB1ngkWTUMhBAsAAGzxTLDwcbgpAADWeShYxG9jJAsAAKzxULBgKAQAANs8GCwsFwIAQDvmnWDR8EnosQAAwB7PBAuHoRAAAKzzTLBwL0IWs1wIAADtmGeCBZM3AQCwz0PBIn5LsAAAwB7vBAsuQgYAgHXeCRYMhQAAYJ2HgkX8ljNvAgBgj3eCBUMhAABY551gwVAIAADWeShYxG8ZCgEAwB4PBQuGQgAAsM2DwYJkAQCALd4JFlyEDAAA6zwTLPwMhQAAYJ1nggVXNwUAwD7PBIvGo0KidFkAAGCNZ4KFvyFZ0GEBAIA9ngkWHBUCAIB9xxQsZs6cKcdxdMMNN6SonNZzGAoBAMC6VgeLNWvW6MEHH9SIESNSWU+rMRQCAIB9rQoWNTU1mjp1qh5++GF17tw51TW1CkMhAADY16pgUVpaqsmTJ2vixIlf2TYSiSgcDics6eAOhRAsAACwJpDsC+bOnas33nhDa9asaVH78vJy3XHHHUkXlqzGE2QZIxlj3PNaAACAzEmqx6KiokLXX3+9/vKXvygnJ6dFrykrK1NVVZW7VFRUtKrQr+I7KEjQaQEAgB1J9VisW7dOu3fv1mmnneaui0ajWr58ue6//35FIhH5/f6E1wSDQQWDwdRUexQHB4uoMfKJHgsAADItqWAxYcIEbdy4MWHdVVddpSFDhujmm28+LFRkku+gvhcmcAIAYEdSwSIvL0/Dhw9PWNexY0d17dr1sPWZxlAIAAD2ee7MmxInyQIAwJakjwo51NKlS1NQxrFjKAQAAPs82WNBhwUAAHZ4M1iQLAAAsMJDwaLpPkMhAADY4Zlg4TiOe1pvOiwAALDDM8FC4kJkAADY5rFgEb8lWAAAYIfHgkVjj4XlQgAAaKe8GSxIFgAAWOGxYBG/ZSgEAAA7vBUsfAyFAABgk7eCBUeFAABglceCRfyWORYAANjhqWDhZygEAACrPBUsHIZCAACwylPBonEoJEqXBQAAVngqWPgbeizosAAAwA5PBQuGQgAAsMtTwcLX8GmiBAsAAKzwVLBoGgohWAAAYIOnggUXIQMAwC5PBQuHo0IAALDKU8Gi6QRZBAsAAGzwVLDwcbgpAABWeSpYNB5uylAIAAB2eCpY+Bs+DUMhAADY4algwVAIAAB2eSpYMBQCAIBdngoW/obDTRkKAQDADk8FC06QBQCAXR4NFiQLAABs8Faw4KgQAACs8lawYCgEAACrvBksSBYAAFiRVLCYNWuWRowYofz8fOXn56ukpEQLFixIV21J83GtEAAArEoqWPTu3VszZ87UunXrtHbtWn3jG9/QJZdcorfffjtd9SXF5x5uarcOAADaq0AyjS+66KKExzNmzNCsWbO0atUqDRs2rNnXRCIRRSIR93E4HG5FmS3DUAgAAHa1eo5FNBrV3LlztXfvXpWUlByxXXl5uUKhkLsUFxe3dpNficNNAQCwK+lgsXHjRnXq1EnBYFDXXHON5s2bp6FDhx6xfVlZmaqqqtyloqLimAo+GoZCAACwK6mhEEkaPHiwNmzYoKqqKj377LOaNm2ali1bdsRwEQwGFQwGj7nQlmjssYjSYwEAgBVJB4vs7GwNGDBAkjR69GitWbNG9957rx588MGUF5esxhNkGYIFAABWHPN5LGKxWMLkTJuYvAkAgF1J9ViUlZVp0qRJ6tOnj6qrq/Xkk09q6dKleumll9JVX1L8DZMs6gkWAABYkVSw2L17t6644grt2LFDoVBII0aM0EsvvaRzzz03XfUlJdAwFkKwAADAjqSCxaOPPpquOlIiy9/QYxGNWa4EAID2yVPXCgk0BIu6KD0WAADY4K1g4Q6F0GMBAIANngoWTUMh9FgAAGCDp4JFwB//OAyFAABgh6eCRZavcY4FQyEAANjgqWDR2GPBHAsAAOzwVLDIYigEAACrPBYsOI8FAAA2eSpYBBrnWHDmTQAArPBWsGicY0GPBQAAVngqWHAeCwAA7PJYsIh/nFp6LAAAsMJTwSInyy9JOlAXtVwJAADtk6eCRW52PFjsqyVYAABgg6eCRYeGHov9BAsAAKzwVrDIDkiSvtxXa7kSAADaJ08Fi6KCHEnSl/vqtK+23nI1AAC0P54KFl06ZisnK/6RPq+m1wIAgEzzVLBwHEfd84KSpM9qDliuBgCA9sdTwUKSunVqCBb0WAAAkHGeCxbdG4NFTcRyJQAAtD/eCxaNQyHVBAsAADLNc8GiaSiEYAEAQKZ5Llg09lh8zlAIAAAZ59lgQY8FAACZR7AAAAAp471gcdBRIcYYy9UAANC+eC5YNE7erK2PqTrCab0BAMgkzwWL3Gy/8oLxi5HtqOTsmwAAZJLngoUkndwrX5L05j8r7RYCAEA748lgMfyEkCTpZ8++ZbkSAADaF08Giz17m44IicWYwAkAQKZ4MlhcP3GQe7/iy30WKwEAoH1JKliUl5dr7NixysvLU2Fhob71rW9py5Yt6aqt1fp16+jef3L1douVAADQviQVLJYtW6bS0lKtWrVKixYtUl1dnc477zzt3bs3XfW1WuMEzgeXfWS5EgAA2o9AMo0XLlyY8Pixxx5TYWGh1q1bp7PPPrvZ10QiEUUiTXMewuFwK8pM3sSTC/Xujvi2/vfNT3XRyKKMbBcAgPbsmOZYVFVVSZK6dOlyxDbl5eUKhULuUlxcfCybbLHrJwx07/94zno9+TpDIgAApJtjWnne61gsposvvliVlZVasWLFEds112NRXFysqqoq5efnt2bTLbbpkypdeF9TbdNKTtTtFw+T4zhp3S4AAF4TDocVCoW+8vc7qaGQg5WWlmrTpk1HDRWSFAwGFQwGW7uZY9J4PotGj6/8WI+v/FgP/NtpGtu3swrzc6zUBQCAV7Wqx+K6667Tc889p+XLl6tfv35JvbaliSdVjDH6/v+s1cvv7m72+Zu/OUQ/PLu/fD56MQAAOJKW/n4nNcfCGKPrrrtO8+bN05IlS5IOFTY4jqNHpo3VjQed2+Jgv1m4Wf1veVFRTqQFAMAxS6rH4tprr9WTTz6p5557ToMHD3bXh0Ih5ebmtug9Mt1jcaiqfXUa+V9/b/a56ycM1I3nNh9AAABoz1r6+51UsDjSpMfZs2fryiuvTGlh6fbl3lp9e9Zr+ujz5s/B8dQPvqZx/btmuCoAAI5PaQkWqXC8BIuD3fjUBs1b/0mzz/1i8sn697P6Z7giAACOLwSLJG3eGdY373n1iM+PPrGznv5hifxM8gQAtENpmbzpZUN65mvbzMn643dGN/v8uo+/1Em3vKjKfbUZrgwAgLaDHosjqI/G9O//s1ZLt3x22HM/OLu/brngZAtVAQBgB0MhKfSvD67U6q17Dlu/tfwCzuIJAGgXGApJoad/WKKNt5+nMwd0S1jfr+xF/XnlNjtFAQBwHKLHIknL3vtM0/60+rD19F4AALyMHos0OWdQd30wY9Jh6/uVvajd1QcsVAQAwPGDYNEKAb9P22ZO1qs/+3rC+tNnLNYDr3xgqSoAAOwjWByD4i4dDuu9uOulLfqkcr+ligAAsItgcYwCfp+2ll+gScN7uuvGz1zCRc0AAO0SwSIFHMfRrO+M1vgBTdcWmfLQKosVAQBgB8Eihf589Tj3/upte+i1AAC0OwSLFPL5HD1zTYn7+IyZiy1WAwBA5hEsUmxs3y7q372jJGlXOKJZSz+0XBEAAJlDsEiDxTed497/zcLNOlAXtVgNAACZQ7BIA8dxdP2Ege7jIbcuVIZPcAoAgBUEizS58dxBCuVmuY9LypdYrAYAgMwgWKTR+lvPde/vDB/ggmUAAM8jWKSRz+doxc1Np/2+9bm3tb+W+RYAAO8iWKRZ784dNKpPgfv4e4+vsVcMAABpRrDIgHnXjnfvv/bhF9r2+V6L1QAAkD4Eiwx56/bz3Pv/8rul9goBACCNCBYZkp+TpUtHneA+fuTVjyxWAwBAehAsMuj/XX6qe//OF97l3BYAAM8hWGTY41ef7t6f/Y9t9goBACANCBYZds6g7soOxHf7fz3/jmJcARUA4CEECwtevrHpWiI/n7/JYiUAAKQWwcKCPl07uPfnrN7OXAsAgGcQLCx56gdfc+9/ncNPAQAeQbCwZFz/rurTJd5zse2LfVq4aafligAAOHYEC4uW/6zpOiLXPLGOiZwAgDaPYGHZo9PGuPcf5qRZAIA2LulgsXz5cl100UUqKiqS4ziaP39+GspqPyac3EO9O+dKksoXbNau8AHLFQEA0HpJB4u9e/dq5MiReuCBB9JRT7v0yEG9Fj979i2LlQAAcGySDhaTJk3SnXfeqUsvvTQd9bRLQ3rm68aJgyRJy977TH9euc1uQQAAtFLa51hEIhGFw+GEBYf74Tn93fu3Pve2aiL1FqsBAKB10h4sysvLFQqF3KW4uDjdm2yTcrL8euSKpiGR6+est1gNAACtk/ZgUVZWpqqqKnepqKhI9ybbrIlDe2jiyT0kSYs379b89Z9YrggAgOSkPVgEg0Hl5+cnLDiyh7472r1/w1MbOEoEANCmcB6L44zP5+iZa0rcx1MfeV310ZjFigAAaLmkg0VNTY02bNigDRs2SJK2bt2qDRs2aPv27amurd0a27eLvn1ab0nSB7trdPZvX7FcEQAALZN0sFi7dq1GjRqlUaNGSZJuuukmjRo1SrfddlvKi2vP/vtfR+q8ofH5Fp9WHdA1f15nuSIAAL6aYzJ8ze5wOKxQKKSqqirmW3yFSH1Ug3+x0H38t+vGa0TvAnsFAQDarZb+fjPH4jgWDPi18Iaz3McX3/8P7a+NWqwIAICjI1gc54b0zNfUcX3cx/e/8r7FagAAODqCRRsw49JT3Puzln6oDz+rsVgNAABHRrBoIzb/6ps65YSQYkaa8N/LVFvPIagAgOMPwaKNyMnyq/z/NPVcDPrFAovVAADQPIJFGzL8hJAGFHZyH7+7gwu6AQCOLwSLNuZ/rzvTvX/LvI1cBRUAcFwhWLQxudl+PVc6XpK0fnulvvvo6wofqLNcFQAAcQSLNmhkcYGeKx2vUG6W1m+v1Ijb/66POFIEAHAcIFi0USOLC/SHqae5j7/x38v0we5qixUBAECwaNPGD+imK8/o6z6eePdybfqkyl5BAIB2j2DRxt1+8TAN6ZnnPr7wvhV651OOFgEA2EGw8ICFN5ytu/7vCPfxBb9/Vf/28CqLFQEA2iuChUdcNqZYz/+46VDU1z78Qn2nv6C6KGfoBABkDsHCQ4afENLLN52dsG7gzxdo/vpPLFUEAGhvCBYeM6AwTxtvPy9h3Q1PbVDf6S8oFjOWqgIAtBcECw/Ky8nStpmTdc/lpyas73/Li/rL6x/bKQoA0C44xpiM/jc2HA4rFAqpqqpK+fn5mdx0u1QTqdfwX7502Pp/TP+GTijItVARAKAtaunvNz0WHtcpGGi292L8zCXqO/0F7a+N2ikMAOBJBIt24lujTtD7Mybpa/27JKw/+baFuuJPqwkYAICUYCikHdqzt1an/WpRs8+tKpugnqGcDFcEADjetfT3m2DRji177zNN+9PqZp97+aazdVL3TnIcJ8NVAQCORwQLtNimT6p04X0rmn2uuEuuHp02VoN65DX7PACgfSBYIGlvVlTqkgf+cdQ2108YqBsmDqQnAwDaGYIFWi0WM/rftz7VY69t0/rtlc226dwhS7lZfj165Vid3IvvEQC8jmCBlDDGaNE7u/SDP6/7yrah3Cwt+ck56topmIHKAACZRLBAWvzjg88144V39c6Or740+40TB+nrQ7prRO+C9BcGAEgrggXSLhozWrhpp0qffCOp102fNET/Mri7BhXmyedjrgYAtAUEC2RcXTSmNVv36N8eeT2p1w0rypcx8cAxtChf3RhKAYDjDsEC1lXtr9Pu8AH9YemHmpfEpdtDuVmq2l/nPu4VytHFpxbptD6dlRcMaEzfLsoOcNJYAMgkggWOW7X1MT386ke666UtKXvPbp2ydWLXjsr2+/S1/l115sBuCuVmqX+3jjKS/Ay5AMAxIVigTarcV6sPP9urG55ar4o9+9O2nW8O66mCDlkqzAvqn5X7NbJ3gYIBn4oKclUXjenkXvnqkZ9DIAGABgQLeI4xRlX76/TOjrB2hyN6b1e1nl77T31eE1GvUI52VB3ISB1nnNRVr334hTp3yNK4fl3VKSeg4s4dVBeNKSfLpy/21mps3y7aXxtVKDdL3fKC6pmfI58jhTpkKeDzEVgAtDlpDRYPPPCA7rrrLu3cuVMjR47Ufffdp9NPPz2lhQHHwhijSH1M7++q0cvv7lKW31HHYEALNu7U6m17NPHkQn1WU6vd4QOqrY9Jkr7YW5ux+nKy4nNE8nOyFMzyKRaTPqmM99CcWlygHvlBOXL0eU1Eudl+9QrlqEN2QF/srVXfrh0U8Pnkc6Qv99Vp886wLjill7L8jjoFs9R4UtSCDlnKyfIr2+9TNGaUHfApO+BTwOfI73PcgBPwOfL7nYT1PkecXRVAgrQFi6eeekpXXHGF/vjHP2rcuHG655579Mwzz2jLli0qLCxMWWGALbGY0Z59tfpwd438Pkcvvb1Tm3dWq0+XDqqPGhV0yNLjK7fpQF1MQ3rmaWhRvj6vqVUsZrTig88lxed8fLmvTtFYRjsEU6oxaORm+9UxOyDHiV8ZN5SbpS4dsxN6XWLGqGN2QNkBn2LGKBjwy+9zdKAuqvycLDfA+Bpupfi8F78bZpyG94kfxpyb7deOyv06sWtH+X2OG3Tc93EcdwJvwOco4I/fjzUEqJwsf8P2pH21UfkdRznZ8ZDlSIoaI7/juOEp4I+/pzHGrcnX8FxjvvI58XX1sZjyc7IU8DuqjxrVx4yCDYGt8f2y/I5iJh5wY0YyMsry+eT3O27NPif+WaIm/lzjtqIx435W4HiStmAxbtw4jR07Vvfff78kKRaLqbi4WD/+8Y81ffr0lBUGeM2BuqjCB+oU8PlUc6BeRsbtJak+UK/56z/RyN4h1UTqlZ+bJcdxVH2gTpG6mBxHqjlQr721UQUafrAdR9pRdUAVe/apuEt8KKYuGv/rHI0Z1UTqVVsfa1gfc3+0ozGj+lhM0Zhx2+P45TiS34mHssaoYYzcYOU0tHEcxw1B8YDld9+j8Z95x4mHNJ/TFNZ8PslR/LEUn1wd8PsU8DuKxox8jqMsv6PGrR/ckdUYpJyD1ru3cg663/RhHB3a7gjvo6YGh7/GOer2Dq7v4Nc6B60/dHs66PO1dHs6wmc42vZ0yOdr2facQ54/fD8dur2fnDdIeTlZSqWW/n4HknnT2tparVu3TmVlZe46n8+niRMnauXKlc2+JhKJKBKJJBQGtEc5WX7lZMX/se/SMVuSdGLXju7z5wzqbqWuWCz+v+6DA0fTY6O9kXrVROolxSfXHqiLqUO2X7GGH6v6qFFtNCZj4ucyMQ0Bpi4WU+W+Ovl9jrL8PkVjMUVjUn00PvQUNcbddn3MuP8419bHFD5Qp3219coLZilSH1XA75PfcRSpj7q11UWNHEfudmPGyJh478mBuqiiRpKJB6zGHo76qJGRaeidiNdgTNPn9fscN3zFew7i7+84cnsg6qJG++uibu+Gr6HmaMP2U80YqT7+wRLW1zbsxyOrT30xaDOu/fpJKQ8WLZVUsPj8888VjUbVo0ePhPU9evTQ5s2bm31NeXm57rjjjtZXCCCtfD5H2W63u/+obdE8Y+LDF7GGMFIfaxxqiT/fOLxR1xAGYsYk9Bi5r29IJr6G+6ZhGCUaM26wafzfb119TKbhtQ0ZSlJT+Di4Nyr+v9n4e7rv2xDCYsYkDNvEh2Ditfl98QDWGAabthOvSw2PG7fU2DNycCkJ7dx1h7Q9pP7Edl+9PR3Szr1/lO3JJL5nS7ang96jaRvmkM+QuL2DX9rS7TW9xiS0Pdr76JDP1zE7qZ/3lEr7lsvKynTTTTe5j8PhsIqLi9O9WQDImMYf+8ZT1Gf5E/vkA/7GeSUEN3hfUsGiW7du8vv92rVrV8L6Xbt2qWfPns2+JhgMKhjkFM0AALQHSZ0XOTs7W6NHj9bixYvddbFYTIsXL1ZJSUnKiwMAAG1L0kMhN910k6ZNm6YxY8bo9NNP1z333KO9e/fqqquuSkd9AACgDUk6WFx++eX67LPPdNttt2nnzp069dRTtXDhwsMmdAIAgPaHU3oDAICv1NLfb649DQAAUoZgAQAAUoZgAQAAUoZgAQAAUoZgAQAAUoZgAQAAUoZgAQAAUoZgAQAAUibj11VtPB9XOBzO9KYBAEArNf5uf9V5NTMeLKqrqyWJS6cDANAGVVdXKxQKHfH5jJ/SOxaL6dNPP1VeXp4cx0nZ+4bDYRUXF6uiooJThVvA/reL/W8X+98u9n9mGGNUXV2toqIi+XxHnkmR8R4Ln8+n3r17p+398/Pz+YNlEfvfLva/Xex/u9j/6Xe0nopGTN4EAAApQ7AAAAAp45lgEQwG9ctf/lLBYNB2Ke0S+98u9r9d7H+72P/Hl4xP3gQAAN7lmR4LAABgH8ECAACkDMECAACkDMECAACkDMECAACkjGeCxQMPPKC+ffsqJydH48aN0+rVq22X1OYsX75cF110kYqKiuQ4jubPn5/wvDFGt912m3r16qXc3FxNnDhR77//fkKbPXv2aOrUqcrPz1dBQYG+973vqaamJqHNW2+9pbPOOks5OTkqLi7Wb3/723R/tDahvLxcY8eOVV5engoLC/Wtb31LW7ZsSWhz4MABlZaWqmvXrurUqZO+/e1va9euXQlttm/frsmTJ6tDhw4qLCzUT3/6U9XX1ye0Wbp0qU477TQFg0ENGDBAjz32WLo/3nFv1qxZGjFihHv2xpKSEi1YsMB9nn2fWTNnzpTjOLrhhhvcdXwHbYTxgLlz55rs7Gzzpz/9ybz99tvm+9//vikoKDC7du2yXVqb8uKLL5qf//zn5q9//auRZObNm5fw/MyZM00oFDLz5883b775prn44otNv379zP79+9023/zmN83IkSPNqlWrzKuvvmoGDBhgpkyZ4j5fVVVlevToYaZOnWo2bdpk5syZY3Jzc82DDz6YqY953Dr//PPN7NmzzaZNm8yGDRvMBRdcYPr06WNqamrcNtdcc40pLi42ixcvNmvXrjVf+9rXzBlnnOE+X19fb4YPH24mTpxo1q9fb1588UXTrVs3U1ZW5rb56KOPTIcOHcxNN91k3nnnHXPfffcZv99vFi5cmNHPe7z529/+Zl544QXz3nvvmS1btphbbrnFZGVlmU2bNhlj2PeZtHr1atO3b18zYsQIc/3117vr+Q7aBk8Ei9NPP92Ulpa6j6PRqCkqKjLl5eUWq2rbDg0WsVjM9OzZ09x1113uusrKShMMBs2cOXOMMca88847RpJZs2aN22bBggXGcRzzySefGGOM+cMf/mA6d+5sIpGI2+bmm282gwcPTvMnant2795tJJlly5YZY+L7OysryzzzzDNum3fffddIMitXrjTGxMOhz+czO3fudNvMmjXL5Ofnu/v8Zz/7mRk2bFjCti6//HJz/vnnp/sjtTmdO3c2jzzyCPs+g6qrq83AgQPNokWLzDnnnOMGC76DtqPND4XU1tZq3bp1mjhxorvO5/Np4sSJWrlypcXKvGXr1q3auXNnwn4OhUIaN26cu59XrlypgoICjRkzxm0zceJE+Xw+vf76626bs88+W9nZ2W6b888/X1u2bNGXX36ZoU/TNlRVVUmSunTpIklat26d6urqEr6DIUOGqE+fPgnfwSmnnKIePXq4bc4//3yFw2G9/fbbbpuD36OxDX9fmkSjUc2dO1d79+5VSUkJ+z6DSktLNXny5MP2E99B25Hxq5um2ueff65oNJrwB0mSevTooc2bN1uqynt27twpSc3u58bndu7cqcLCwoTnA4GAunTpktCmX79+h71H43OdO3dOS/1tTSwW0w033KDx48dr+PDhkuL7Jzs7WwUFBQltD/0OmvuOGp87WptwOKz9+/crNzc3HR+pTdi4caNKSkp04MABderUSfPmzdPQoUO1YcMG9n0GzJ07V2+88YbWrFlz2HP8+W872nywALyotLRUmzZt0ooVK2yX0q4MHjxYGzZsUFVVlZ599llNmzZNy5Yts11Wu1BRUaHrr79eixYtUk5Oju1ycAza/FBIt27d5Pf7D5sZvGvXLvXs2dNSVd7TuC+Ptp979uyp3bt3JzxfX1+vPXv2JLRp7j0O3kZ7d9111+n555/XK6+8ot69e7vre/bsqdraWlVWVia0P/Q7+Kr9e6Q2+fn57f5/a9nZ2RowYIBGjx6t8vJyjRw5Uvfeey/7PgPWrVun3bt367TTTlMgEFAgENCyZcv0+9//XoFAQD169OA7aCPafLDIzs7W6NGjtXjxYnddLBbT4sWLVVJSYrEyb+nXr5969uyZsJ/D4bBef/11dz+XlJSosrJS69atc9ssWbJEsVhM48aNc9ssX75cdXV1bptFixZp8ODB7X4YxBij6667TvPmzdOSJUsOGzIaPXq0srKyEr6DLVu2aPv27QnfwcaNGxMC3qJFi5Sfn6+hQ4e6bQ5+j8Y2/H05XCwWUyQSYd9nwIQJE7Rx40Zt2LDBXcaMGaOpU6e69/kO2gjbs0dTYe7cuSYYDJrHHnvMvPPOO+YHP/iBKSgoSJgZjK9WXV1t1q9fb9avX28kmbvvvtusX7/efPzxx8aY+OGmBQUF5rnnnjNvvfWWueSSS5o93HTUqFHm9ddfNytWrDADBw5MONy0srLS9OjRw3z3u981mzZtMnPnzjUdOnTgcFNjzI9+9CMTCoXM0qVLzY4dO9xl3759bptrrrnG9OnTxyxZssSsXbvWlJSUmJKSEvf5xsPtzjvvPLNhwwazcOFC071792YPt/vpT39q3n33XfPAAw9wuJ0xZvr06WbZsmVm69at5q233jLTp083juOYv//978YY9r0NBx8VYgzfQVvhiWBhjDH33Xef6dOnj8nOzjann366WbVqle2S2pxXXnnFSDpsmTZtmjEmfsjprbfeanr06GGCwaCZMGGC2bJlS8J7fPHFF2bKlCmmU6dOJj8/31x11VWmuro6oc2bb75pzjzzTBMMBs0JJ5xgZs6cmamPeFxrbt9LMrNnz3bb7N+/31x77bWmc+fOpkOHDubSSy81O3bsSHifbdu2mUmTJpnc3FzTrVs385Of/MTU1dUltHnllVfMqaeearKzs03//v0TttFeXX311ebEE0802dnZpnv37mbChAluqDCGfW/DocGC76BtcIwxxk5fCQAA8Jo2P8cCAAAcPwgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZQgWAAAgZf4/ltdNHMI9R7kAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# stage 2: train to 24 bits\n",
    "randkey, trainkey = jax.random.split(randkey)\n",
    "model = train_stage(trainkey, model, optimizer, batch_size,\n",
    "                    iterations=16384, min_val=0, max_val=(1 << 24) - 1)\n",
    "test_fizzbuzz(model, [i for i in range(int(1e2), int(1e2) + 25)])\n",
    "test_fizzbuzz(model, [i for i in range(int(1e6), int(1e6) + 25)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "63728fdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b67500c28f10414090e446f488dc3394",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/32768 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "200  fizz: 0.0000, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "201  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 1, neither: 0.0000\n",
      "202  fizz: 0.0000, fizzmod: 1, buzz: 0.0013, buzzmod: 2, neither: 0.9997\n",
      "203  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "204  fizz: 1.0000, fizzmod: 0, buzz: 0.0022, buzzmod: 4, neither: 0.0000\n",
      "205  fizz: 0.0000, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "206  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 1, neither: 1.0000\n",
      "207  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 2, neither: 0.0000\n",
      "208  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "209  fizz: 0.0000, fizzmod: 2, buzz: 0.0007, buzzmod: 4, neither: 0.9951\n",
      "210  fizz: 1.0000, fizzmod: 0, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "211  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 1, neither: 0.9997\n",
      "212  fizz: 0.0000, fizzmod: 2, buzz: 0.0008, buzzmod: 2, neither: 0.9842\n",
      "213  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 3, neither: 0.0000\n",
      "214  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 4, neither: 0.9999\n",
      "215  fizz: 0.0000, fizzmod: 2, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "216  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 1, neither: 0.0000\n",
      "217  fizz: 0.0000, fizzmod: 1, buzz: 0.0182, buzzmod: 2, neither: 0.9998\n",
      "218  fizz: 0.0001, fizzmod: 2, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "219  fizz: 1.0000, fizzmod: 0, buzz: 0.0001, buzzmod: 4, neither: 0.0000\n",
      "220  fizz: 0.0000, fizzmod: 1, buzz: 1.0000, buzzmod: 0, neither: 0.0000\n",
      "221  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 1, neither: 1.0000\n",
      "222  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 2, neither: 0.0000\n",
      "223  fizz: 0.0000, fizzmod: 1, buzz: 0.0000, buzzmod: 3, neither: 1.0000\n",
      "224  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 4, neither: 0.9999\n",
      "1000000000  fizz: 0.0002, fizzmod: 1, buzz: 0.9998, buzzmod: 0, neither: 0.0016\n",
      "1000000001  fizz: 0.0001, fizzmod: 2, buzz: 0.0029, buzzmod: 1, neither: 0.9992\n",
      "1000000002  fizz: 0.9942, fizzmod: 0, buzz: 0.0000, buzzmod: 2, neither: 0.0374\n",
      "1000000003  fizz: 0.0001, fizzmod: 1, buzz: 0.0001, buzzmod: 3, neither: 0.9997\n",
      "1000000004  fizz: 0.0001, fizzmod: 2, buzz: 0.0003, buzzmod: 4, neither: 0.9969\n",
      "1000000005  fizz: 1.0000, fizzmod: 0, buzz: 0.9991, buzzmod: 0, neither: 0.0000\n",
      "1000000006  fizz: 0.0003, fizzmod: 1, buzz: 0.0190, buzzmod: 1, neither: 0.9696\n",
      "1000000007  fizz: 0.0000, fizzmod: 2, buzz: 0.0000, buzzmod: 2, neither: 1.0000\n",
      "1000000008  fizz: 0.9991, fizzmod: 0, buzz: 0.0342, buzzmod: 3, neither: 0.0000\n",
      "1000000009  fizz: 0.0002, fizzmod: 1, buzz: 0.0003, buzzmod: 4, neither: 0.9999\n",
      "1000000010  fizz: 0.0145, fizzmod: 2, buzz: 0.9907, buzzmod: 0, neither: 0.0010\n",
      "1000000011  fizz: 0.9938, fizzmod: 0, buzz: 0.0046, buzzmod: 1, neither: 0.0015\n",
      "1000000012  fizz: 0.0002, fizzmod: 1, buzz: 0.0001, buzzmod: 2, neither: 1.0000\n",
      "1000000013  fizz: 0.0001, fizzmod: 2, buzz: 0.0238, buzzmod: 3, neither: 0.9997\n",
      "1000000014  fizz: 0.9980, fizzmod: 0, buzz: 0.0003, buzzmod: 4, neither: 0.0040\n",
      "1000000015  fizz: 0.0003, fizzmod: 1, buzz: 0.9948, buzzmod: 0, neither: 0.0747\n",
      "1000000016  fizz: 0.0002, fizzmod: 2, buzz: 0.0012, buzzmod: 1, neither: 0.9977\n",
      "1000000017  fizz: 1.0000, fizzmod: 0, buzz: 0.0000, buzzmod: 2, neither: 0.0009\n",
      "1000000018  fizz: 0.0002, fizzmod: 1, buzz: 0.0002, buzzmod: 3, neither: 0.9987\n",
      "1000000019  fizz: 0.0001, fizzmod: 2, buzz: 0.0000, buzzmod: 4, neither: 1.0000\n",
      "1000000020  fizz: 1.0000, fizzmod: 0, buzz: 0.9998, buzzmod: 0, neither: 0.0000\n",
      "1000000021  fizz: 0.0001, fizzmod: 1, buzz: 0.0017, buzzmod: 1, neither: 0.9983\n",
      "1000000022  fizz: 0.0001, fizzmod: 2, buzz: 0.0000, buzzmod: 2, neither: 1.0000\n",
      "1000000023  fizz: 1.0000, fizzmod: 0, buzz: 0.0002, buzzmod: 3, neither: 0.0000\n",
      "1000000024  fizz: 0.0004, fizzmod: 1, buzz: 0.0001, buzzmod: 4, neither: 0.9999\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhYAAAGdCAYAAABO2DpVAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAsUElEQVR4nO3deXxU9b3/8feZmWSSQBa2JCBhExVlUZCliDsoRaxL+7DqpRW1y08b68KtLdirt95WQ+vj3p+tUqreFnuvC9X+1KpFKYtCrbJaEERZBCWCYTWZBMJkme/vj2ROMuyTzOSbnHk9H4959Jwz35nzmTmVeef7/Z5zHGOMEQAAQAL4bBcAAAC8g2ABAAAShmABAAAShmABAAAShmABAAAShmABAAAShmABAAAShmABAAASJtDWO4xEItq5c6eys7PlOE5b7x4AALSAMUaVlZXq1auXfL5j90u0ebDYuXOnioqK2nq3AAAgAUpLS9W7d+9jPt/mwSI7O1tSQ2E5OTltvXsAANACoVBIRUVF7u/4sbR5sIgOf+Tk5BAsAADoYE40jYHJmwAAIGEIFgAAIGEIFgAAIGEIFgAAIGEIFgAAIGEIFgAAIGEIFgAAIGEIFgAAIGHiChb9+vWT4zhHPIqLi5NVHwAA6EDiuvLmypUrVV9f766vX79el112ma677rqEFwYAADqeuIJFjx49YtZnzpypU089VRdddFFCiwIAAB1Ti+8VUlNTo2eeeUbTpk077nXDw+GwwuGwux4KhVq6SwAA0M61ePLmK6+8ovLyct18883HbVdSUqLc3Fz3kaxbpv/n3zbqZ69+qLKKQ0l5fwAAcGKOMca05IUTJ05Uenq6XnvtteO2O1qPRVFRkSoqKhJ6d9NRDy3Unsqw5t15gc7qxV1TAQBIpFAopNzc3BP+frdoKOSzzz7TwoUL9dJLL52wbTAYVDAYbMluWsSoRTkJAAAkQIuGQubMmaP8/HxNnjw50fW0WHSWR8v6XwAAQCLEHSwikYjmzJmjqVOnKhBo8dzPhDvO/FEAANBG4g4WCxcu1Pbt23Xrrbcmox4AANCBxd3lcPnll6uF8z2TyhFdFgAA2Oa5e4W0w8wDAEDK8EywiM6x4KwQAADs8U6wsF0AAADwTrCIYigEAAB7PBMsjne/EgAA0DY8Eyyi6LAAAMAe7wULxkIAALDGM8GCkRAAAOzzTLCIor8CAAB7PBMs6LEAAMA+zwSLKKZYAABgj2eCBfcKAQDAPs8EiyZ0WQAAYItngoV7rxByBQAA1ngnWNguAAAAeCdYRNFhAQCAPZ4JFtwrBAAA+zwTLKKYYwEAgD2eCRbR/gruFQIAgD2eCRbM3gQAwD7vBItG9FcAAGCPZ4IFHRYAANjnmWARxRQLAADs8UywiJ5uahgMAQDAGu8EC9sFAAAA7wQLFx0WAABY45lgwYU3AQCwzzPBIooOCwAA7PFMsHAaZ1lwVggAAPZ4J1gwFAIAgHWeCRZRnG4KAIA9ngsWAADAHs8FC+ZYAABgj2eChcMkCwAArPNMsIiiwwIAAHs8Eyyi/RWGsRAAAKzxTrBgJAQAAOviDhY7duzQt771LXXr1k2ZmZkaOnSoVq1alYzaWoT+CgAA7AnE0/jLL7/UuHHjdMkll+iNN95Qjx49tHnzZnXp0iVZ9Z00eiwAALAvrmDxy1/+UkVFRZozZ467rX///gkvqlXosgAAwJq4hkJeffVVjRw5Utddd53y8/M1fPhwPfXUU8d9TTgcVigUinkkg3uvEJIFAADWxBUstm7dqtmzZ+u0007T/Pnzdfvtt+vOO+/UH//4x2O+pqSkRLm5ue6jqKio1UUfDUMhAADY55g4zs9MT0/XyJEj9e6777rb7rzzTq1cuVLvvffeUV8TDocVDofd9VAopKKiIlVUVCgnJ6cVpce66vF39MHnFfr91JEaf2ZBwt4XAAA0/H7n5uae8Pc7rh6Lnj176qyzzorZduaZZ2r79u3HfE0wGFROTk7MIxnosAAAwL64gsW4ceO0cePGmG2bNm1S3759E1pUa3B9LAAA7IkrWNxzzz1atmyZHn74YW3ZskXPPfecnnzySRUXFyervpPnRCdvAgAAW+IKFqNGjdLLL7+s559/XkOGDNHPf/5zPfroo5oyZUqy6jtpDIUAAGBfXNexkKQrr7xSV155ZTJqSQjuFQIAgD3cKwQAACSMZ4JFFP0VAADY45lg0XTbdKtlAACQ0rwTLBgLAQDAOs8EiyZ0WQAAYItnggX9FQAA2OeZYBHFHAsAAOzxTLBgigUAAPZ5JlhE0WEBAIA9ngkWTuMsC4ZCAACwxzPBgtmbAADY551g0cgwGAIAgDWeCRZ0WAAAYJ9ngkUUcywAALDHM8EieropuQIAAHu8EywYDAEAwDrPBIsow1gIAADWeCZYcOVNAADs80ywAAAA9nkmWLiTNxkJAQDAGu8ECyZvAgBgnWeCRRRX3gQAwB7PBAsmbwIAYJ9ngkUUcywAALCHYAEAABLGM8HCYSwEAADrPBMsouiwAADAHs8EC/orAACwzzPBIop7hQAAYI9nggVTLAAAsM8zwSKK/goAAOzxTLBwOyxIFgAAWOOdYMFYCAAA1nkmWERxrxAAAOzxTLCgvwIAAPs8EyyiONsUAAB7PBMsolMsyBUAANgTV7D42c9+JsdxYh6DBg1KVm1xYjAEAADbAvG+YPDgwVq4cGHTGwTifoukYigEAAB74k4FgUBAhYWFyailVTjbFAAA++KeY7F582b16tVLAwYM0JQpU7R9+/bjtg+HwwqFQjGPZOJ0UwAA7IkrWIwZM0ZPP/203nzzTc2ePVvbtm3TBRdcoMrKymO+pqSkRLm5ue6jqKio1UUfTbTDgqEQAADsiStYTJo0Sdddd52GDRumiRMnat68eSovL9cLL7xwzNfMmDFDFRUV7qO0tLTVRR8NQyEAANjXqpmXeXl5Ov3007Vly5ZjtgkGgwoGg63ZTVzosAAAwJ5WXceiqqpKn3zyiXr27JmoelrM4XRTAACsiytY/OhHP9KSJUv06aef6t1339W1114rv9+vG2+8MVn1xY9JFgAAWBPXUMjnn3+uG2+8Ufv27VOPHj10/vnna9myZerRo0ey6jtpXHkTAAD74goWc+fOTVYdrcbkTQAA7PPMvUKiGAkBAMAezwQLJm8CAGCfZ4JFlKHLAgAAa7wTLOiwAADAOu8Ei0b0VwAAYI9nggX3CgEAwD7vBAvONwUAwDrPBIsoOiwAALDHM8GC/goAAOzzTLCI4nRTAADs8UywYIoFAAD2eSdY2C4AAAB4J1hEMRICAIA9ngkWnG4KAIB9ngkWUYYTTgEAsMYzwYIrbwIAYJ9nggWzNwEAsM87waIRHRYAANjjmWDh0GUBAIB1ngkWUcyxAADAHs8Ei+jZppwVAgCAPd4JFrYLAAAA3gkWUQyFAABgj2eCBRfeBADAPs8ECwAAYJ9nggWnmwIAYJ9ngkWUYZIFAADWeCZYuKebkisAALDGc8ECAADY45lgEUWHBQAA9ngoWNBlAQCAbR4KFg2YYwEAgD2eCRbcKwQAAPu8EyxsFwAAALwTLKIYCgEAwB7PBAtONwUAwD7PBIsoOiwAALDHM8HCvVcIYyEAAFjTqmAxc+ZMOY6ju+++O0HltBxDIQAA2NfiYLFy5Uo98cQTGjZsWCLraTX6KwAAsKdFwaKqqkpTpkzRU089pS5duiS6phahwwIAAPtaFCyKi4s1efJkTZgw4YRtw+GwQqFQzCOZmGIBAIA9gXhfMHfuXL3//vtauXLlSbUvKSnRgw8+GHdh8XIaJ1lw5U0AAOyJq8eitLRUd911l5599lllZGSc1GtmzJihiooK91FaWtqiQgEAQPsXV4/F6tWrtXv3bo0YMcLdVl9fr6VLl+rxxx9XOByW3++PeU0wGFQwGExMtSeBoRAAAOyJK1iMHz9e69ati9l2yy23aNCgQfrJT35yRKhoS5xuCgCAfXEFi+zsbA0ZMiRmW6dOndStW7cjtttChwUAAPZ478qbAADAmrjPCjnc22+/nYAyEoc5FgAA2OOdHovorUIYDAEAwBrvBAvbBQAAAO8ECxcdFgAAWOOZYMHppgAA2OeZYBFFhwUAAPZ4Jli49wrhtBAAAKzxTrCwXQAAAPBOsIiiwwIAAHu8EyzosgAAwDrvBItGdFgAAGCPZ4JF9F4hDIUAAGCPd4IFQyEAAFjnmWARxb1CAACwxzPBgg4LAADs806wiN7dlA4LAACs8Uyw8PsaPkp9hGQBAIAtngkWAV9Dl0UdwQIAAGs8Eyz8jcGiPhKxXAkAAKnLM8GCHgsAAOzzTrDwM8cCAADbvBMs6LEAAMA6zwQLd45FPcECAABbPBMs6LEAAMA+zwQLzgoBAMA+zwSLgJ8eCwAAbPNMsODKmwAA2OeZYMEcCwAA7PNMsGiaY0GwAADAFs8EC3osAACwzzPBgrNCAACwzzPBItA4ebOOC2QBAGCNZ4IFcywAALDPM8Eieh0LggUAAPZ4Jlj4mbwJAIB1ngkWAYZCAACwzjPBoqnHgrNCAACwxTPBIsAlvQEAsC6uYDF79mwNGzZMOTk5ysnJ0dixY/XGG28kq7a4MMcCAAD74goWvXv31syZM7V69WqtWrVKl156qa6++mp9+OGHyarvpLlzLLiOBQAA1gTiafy1r30tZv2hhx7S7NmztWzZMg0ePDihhcWLHgsAAOyLK1g0V19frxdffFEHDhzQ2LFjj9kuHA4rHA6766FQqKW7PC6uYwEAgH1xT95ct26dOnfurGAwqNtuu00vv/yyzjrrrGO2LykpUW5urvsoKipqVcHHwlkhAADYF3ewOOOMM7RmzRotX75ct99+u6ZOnaoNGzYcs/2MGTNUUVHhPkpLS1tV8LFEzwqJGClCrwUAAFbEPRSSnp6ugQMHSpLOPfdcrVy5Ur/+9a/1xBNPHLV9MBhUMBhsXZUnIdpjIUn1xsgn5zitAQBAMrT6OhaRSCRmDoUtgebBgh4LAACsiKvHYsaMGZo0aZL69OmjyspKPffcc3r77bc1f/78ZNV30pr3WNTUR5SR5rdYDQAAqSmuYLF7927ddNNN+uKLL5Sbm6thw4Zp/vz5uuyyy5JV30kLBpo6X2rqmMAJAIANcQWL3//+98mqo9Ucx1Ew4FO4LqJDtfW2ywEAICV55l4hktzhj0O19FgAAGCDp4JFdDiEHgsAAOzwVLCI9liEmWMBAIAVHgsWDR8nTI8FAABWeCxYNM6xqCNYAABgg7eCRYDJmwAA2OSpYBFMY/ImAAA2eStYBJi8CQCATZ4KFhn0WAAAYJXHggVzLAAAsMlTwYILZAEAYJengkW0x+LdT/ZargQAgNTkqWCx/0CNJKlP106WKwEAIDV5KljkZwclSf/v/c8tVwIAQGryVLD44PMK2yUAAJDSPBUsvn/hAElSmt+xXAkAAKnJU8EiLytNkpSfnWG5EgAAUpOngkWnYECSVM3ppgAAWOGpYJHZeLrpgXCd5UoAAEhNngoW0R6LcF1ENdwvBACANuepYJGbmeYu76kKW6wEAIDU5Klg4fc1nQ1SVlFtsRIAAFKTp4JFc48u3Gy7BAAAUo5ng8V5p3a3XQIAACnHc8FidP+ukqRlW/dZrgQAgNTjuWARbryGRfTUUwAA0HY8FyxuGttPklTFtSwAAGhzngsWvbtkSpLe2bKXa1kAANDGPBcsTmkMFpL09817LFYCAEDq8VywKMxpugHZ9JfWWawEAIDU47lgEfA3faQ9lVx9EwCAtuS5YCFJ379wgLtsjLFYCQAAqcWTwWLN9vKm5dLyY7YDAACJ5clgccPoInf52t++a7ESAABSiyeDxVVn97JdAgAAKcmTwSLg92lEnzx3vbae61kAANAWPBksJOlHl5/hLt/74lqLlQAAkDriChYlJSUaNWqUsrOzlZ+fr2uuuUYbN25MVm2tMvbUbu7yK2t2WqwEAIDUEVewWLJkiYqLi7Vs2TItWLBAtbW1uvzyy3XgwIFk1ddijuPErG/aVWmpEgAAUodjWnGhhz179ig/P19LlizRhRdeeFKvCYVCys3NVUVFhXJyclq665PyyPyPNeutT9z1T2dOTur+AADwqpP9/W7VHIuKigpJUteuXY/ZJhwOKxQKxTzaSvN5FpIUOlTbZvsGACAVtThYRCIR3X333Ro3bpyGDBlyzHYlJSXKzc11H0VFRcdsm2iHD4cM/48FbbZvAABSUYuDRXFxsdavX6+5c+cet92MGTNUUVHhPkpLS1u6yxZ57Mbh7nJ9hMt7AwCQTC0KFnfccYdef/11vfXWW+rdu/dx2waDQeXk5MQ82tLXDrtY1qy3trTp/gEASCVxBQtjjO644w69/PLLWrx4sfr375+suhJqxU/Hu8uPzN/IjckAAEiSuIJFcXGxnnnmGT333HPKzs5WWVmZysrKVF1dnaz6EiI/OyNm/am/b7VUCQAA3hZXsJg9e7YqKip08cUXq2fPnu7jT3/6U7LqS5jpkwa5yw/P+5heCwAAkiDuoZCjPW6++eYklZc43zk/dtim/4x5lioBAMC7PHuvkMOl+X16/Yfnx2yr4+ZkAAAkVMoEC0kackquRvbt4q7fOfefFqsBAMB7UipYSNKLt411l+etK7NYCQAA3pNyweLwq3HuKG/fZ7QAANCRpFywkGKvazFu5mKLlQAA4C0pGSwOv65FuK7eUiUAAHhLSgYLSbpr/Gnu8v2vrLdYCQAA3pGyweKey053l19Y9bnFSgAA8I6UDRaSYk49BQAArZfSwWLKV/q4y4dqmWcBAEBrpXSwuOrsU9zl/7tgk8VKAADwhpQOFn5f0zUtnljKHU8BAGitlA4WAAAgsVI+WDx100h3uT7CrdQBAGiNlA8WI/rkucurP/vSXiEAAHhAygeLTsGAu/zO5j0WKwEAoONL+WCRkeZ3l19Zs9NiJQAAdHwpHywk6baLTpUkbd9/0HIlAAB0bAQLSVef08tdNoYJnAAAtBTBQtLA/M7u8q/mb7RYCQAAHRvBQlKav+lrmP32JxYrAQCgYyNYNPrq4EJ3metZAADQMgSLRsWXDHSXX3qf26gDANASBItGQ07JcZfv/fMHFisBAKDjIlg0chwnZr26htuoAwAQL4JFMxt/8VV3+cwH3rRYCQAAHRPBoplgwH/iRgAA4JgIFof55TeGustrS8vtFQIAQAdEsDjM9aP6uMtXz/qHxUoAAOh4CBYnsHQTdzwFAOBkESyOYum9l7jLN/1hhcVKAADoWAgWR9GnW1bMeil3PQUA4KQQLI4heit1SfruH1dZrAQAgI6DYHEM0ycNcpc37qrUoVoumAUAwIkQLI7j/ivPcpcf+utHFisBAKBjIFgcx63j+rnL/7vsM3uFAADQQRAsjsNxHF1yRg93fc4/tlmsBgCA9i/uYLF06VJ97WtfU69eveQ4jl555ZUklNV+zJoywl1+8LUNFisBAKD9iztYHDhwQGeffbZmzZqVjHranaz0gO0SAADoMOIOFpMmTdIvfvELXXvttcmop136r2+e7S7f+vRKi5UAANC+JX2ORTgcVigUinl0NF8f0dtdXvzxbtVHjMVqAABov5IeLEpKSpSbm+s+ioqKkr3LpPjP65p6LU69b57FSgAAaL+SHixmzJihiooK91FaWprsXSbF10ecErNeXcMFswAAOFzSg0UwGFROTk7MoyNyHEdv3HWBu37mA29arAYAgPaJ61jE4cyesaGIy3wDABAr7mBRVVWlNWvWaM2aNZKkbdu2ac2aNdq+fXuia2uX7plwurs86H56LQAAaC7uYLFq1SoNHz5cw4cPlyRNmzZNw4cP1wMPPJDw4tqjO8cPjFmn1wIAgCZxB4uLL75YxpgjHk8//XQSymt/HMeJufMpvRYAADRhjkUL3HbRqTHrc1ekxjAQAAAnQrBoobd/dLG7PP2ldfqiotpeMQAAtBMEixbq172TumSluetjSxZbrAYAgPaBYNEKq//tspj1BRt2WaoEAID2gWDRCj6fE3ODsu/9zyrV1EUsVgQAgF0Ei1ZqfoMySTr9397gFFQAQMoiWCTAtpIrYtY5BRUAkKoIFgngOI5evWNczLbS/QctVQMAgD0EiwQZ1jtPD1412F2/4FdvWawGAAA7CBYJNPW8fjHrSzbtsVMIAACWECwSbMN/THSXp/5hhcVKAABoewSLBMtKD+jSQfnu+tY9VRarAQCgbREskuD3U0e6y5f+5xKLlQAA0LYIFkngOI5G9+/qrn/rv5dbrAYAgLZDsEiS//3OaHf5nS17dcsc5lsAALyPYJEkwYBfP7i46fbqb23co37T/2qxIgAAko9gkUQ//uqgmCERSeo3/a+KRIyligAASC6CRZK98H/G6ppzesVsG3DfPFXXcD8RAID3ECzawKM3DNff7rkwZtuZD7yplZ/ut1QRAADJQbBoI6cXZOuDn10es+26372n7/5xpaWKAABIPIJFG8rJSDviTqgLP9qtftP/qnrmXQAAPIBg0cYcx9HWh6/Q5WcVxGw/9b55WltabqcoAAASxDHGtOmfyqFQSLm5uaqoqFBOTk5b7rrd2VlerfNmLj5i+/L7xqsgJ8NCRQAAHN3J/n7TY2FRr7zMI4ZGJGnMw4vUb/pftfqzL9XGuQ8AgFYhWFjmOI4+nTlZ148sOuK5b8x+V/1nzNOKbZw9AgDoGBgKaUeMMZry38v17if7jvr89y8coBmTBslxnDauDACQ6k7295tg0U5994+rtPCjXcd8/r0Zl6pnbmYbVgQASGUEC49YsGGXvvc/q475/JXDeuo3NwyXz0cvBgAgeQgWHlNRXauzH/zbcdt874L+uu+KMxkqAQAkHMHCw/65/Utd+9t3j9vmwtN76L9vGqn0APNzAQCtR7BIAZGI0Q1PLTups0b+65tna/KwngoG/G1QGQDAawgWKeZgTZ2uevwf2rK76qRfM6B7J/36huEa2js3iZUBALyAYJHinlz6iR6e93GLX//4vwzXxMGFSvMzlAIAIFjgMFt2V2nCfy1p9fsEfI6e+Pa5GjewuzLSGFYBgFRBsMBxVdfU68mlW/Xn90tVur86oe9d1DVT/3H1EI3o00W5mWkJfW8AgB0EC7RYJGL0j0/26rt/XKVwXSQp+3AcaegpufrmyCJ1DgZ0yaB8ZaT5mFwKAO0UwQJJYYzR2s8rtOrT/Xpi6VbtqQy3eQ19umapa6d0DczvrPNO7aZTe3TWWb1yFGi8SBjX8QCAxCNYwApjjOoiRhvLKvXows1at6Ncu0JtHz5O5LT8zhraO1f7qmq0ZNMeDe6Vo+F98vTMsu365sjeGtY7Tz7HUemXBzW6f1cVdclUj+wMdUr3K9A4odUYI2PEVU8BpISkBotZs2bpkUceUVlZmc4++2w99thjGj16dEILQ2qorY9oY1ml9h+o0aMLN2lXKKx+3bP0jy1HvxFbe5eZ5lenYEB7q2LD1Oj+XbVi237lZAR0xdCeOlBTr9fW7jzi9ZedVaA9lWFt23tA9048Q3sqw1r7ebk++iKkS87IV+WhOn1eXq2bvtJXg3pm628f7tKVw3pq+bb96tetk3pkB+U4TXWk+R2l+X3y+xzO8AHQKkkLFn/6059000036Xe/+53GjBmjRx99VC+++KI2btyo/Pz8hBUGnEhNXUR1kYi27T2gnIw0fbKnSiu27dffN+/Vtr0HVBWuU5+uWRo3sJueX1Fqu9wOp3/3Ttq290DMtnS/TzX1R593c9HpPfRxWUh+x9HOikOaOLhAa0srVJCboYyAT8u37deofl3UKy9TPsdRut+nwtwMvbJmhz7bd1A/vHSgauoj6pQeUFW4ThkBnzLTAwodqtUX5dUa0beLXlu7Uys//VLTLjtd736yV+ed2l2rP/tSVw7rqa17D2jN9nL17ZalEX27KCvdr87BgCQpYow27arSji+rdU5RnspCh3ROUZ47fFYVrlNuZpp2lFcrNzNNuZlpSvP7FAz4VG+MwrUR7a0KqzA3Q5mNZ0PV1EdUfrBWaX6feuVlyBiprt6oLhJRVnpAPp/kc5zGhxQxks+RjJEC/qagZ4xkZOTIUcDf0N5Rwzyk6GuaD+9F/8l2HEf1kYZl/1F6zYwxDAsioZIWLMaMGaNRo0bp8ccflyRFIhEVFRXphz/8oaZPn56wwoD2oK4+Ip/jyEiqPFSriJE+LgtpT2VYxjT8YO34slpGDXM/PtlTpdL9B2Ukzf+wTIdqI0oP+DRpSKHmrftC37tggLLS/fps30G9uPrzI/aXHQyoMlwnSRrRJ0/vby9vy4+LdspxGkKKJDdMZKT5dKi2IeSlB3wNAUSOHCcauhvade2ULp/jKNphFW3TEF4c9/2jGeTw552GjXLcWpqCT7Stmr9Xs/dr/l46/PmjvJea7av569/9ZJ+6dUrXgB6dtK+qRlnBhtAYiTS0T/M7bltjGoJWNJhJDRPSo4IBn0xjHT5fU+gLNAtntfWRmFB2tNHOgM/X8CZGqotElOb3yTS+sc9x5G9837pIxP1em33Ewz577Penw7+7w74Pp7GRE/N+sa/514lnKCcjsWflJSVY1NTUKCsrS3/+8591zTXXuNunTp2q8vJy/eUvfzniNeFwWOFwU7dwKBRSUVERwQJIoOh/xhHT8MMTrqvXwZp6GSPtrKiW33G070BYmWkBdemUpqpDdTpUG9HO8mp17ZSulZ/t17JP9mnikELlZqZp+/6DOrMwR+t3VGj/wRr1yA5qY1ml0v0+fVxWqQPhOg3vk6eFH+3WlcN66vUPvtAPLx2ocF1ETy7dql65GfrGub21bkeFuncOqrqmXn9d94WCAZ/Gn5kvn+MoXBdR987pmv/hLnXvnK5eeZla/emXGtW/qz7de0ADenTWO1v2aNgpeVrxaexl67t3TtfeqhoV5ASPOodncK8chesiCjbeK2dfVY3KQodi2vTtluX+pb91T1PPTGaaX906p6u2PqKauogipuEmgFLDP96d0wMyarjabbPfqyN6c6I/coANK346XvnZGQl9z5MNFoF43nTv3r2qr69XQUFBzPaCggJ9/PHRr/JYUlKiBx98MJ7dAIhT9K8dv9Pw11p6wKfsxr9WCnNP/I/LhLMKjrr9muGnnNT+H/+XpuX7rjjzqG1mHeO1JV8/qV20W82HHA6f0GuMcXsYHMdx/3qVGnoV6iNGAb9P9RGjSMQ09DY09pJVHapTTmaa6iNGRg3ve6i2XrX1DX+mB3yOsoJ+VdfUy+9z3BATHVppmOfjKDsjoJq6SEzQibZp+N+mYGoUbdP8uSOfN2p4MmbdXW78HhpfEN3P0Z5vvq/D379524gxqq0zSgs4WrFtv0b27aq0gE9+x5FRw3fcvI66+oZtvsY/39Mah5gijcNaPsdRXeP3Gh1uqo80Hkdj5PM5qq2LyO/3uT0Gh6trDJGO4zTW2NSzEWnszZTUbMjr8O+wadvhx+Ro30Xz/781PzaN36C7LbqxU3pcP+8JlfQ9z5gxQ9OmTXPXoz0WAOAFzbvMHaepezq6HvA3bfD7mq7TcqJrtnTvHGxVXX27dWrV69ura4f3tl0CTiCuYNG9e3f5/X7t2rUrZvuuXbtUWFh41NcEg0EFg637DwQAAHQMcZ1/lp6ernPPPVeLFi1yt0UiES1atEhjx45NeHEAAKBjiXsoZNq0aZo6dapGjhyp0aNH69FHH9WBAwd0yy23JKM+AADQgcQdLK6//nrt2bNHDzzwgMrKynTOOefozTffPGJCJwAASD1c0hsAAJzQyf5+c41fAACQMAQLAACQMAQLAACQMAQLAACQMAQLAACQMAQLAACQMAQLAACQMAQLAACQMG1+X9Xo9bhCoVBb7xoAALRQ9Hf7RNfVbPNgUVlZKUncOh0AgA6osrJSubm5x3y+zS/pHYlEtHPnTmVnZ8txnIS9bygUUlFRkUpLS7lUeAfBMet4OGYdD8es42mvx8wYo8rKSvXq1Us+37FnUrR5j4XP51Pv3r2T9v45OTnt6kDgxDhmHQ/HrOPhmHU87fGYHa+nIorJmwAAIGEIFgAAIGE8EyyCwaD+/d//XcFg0HYpOEkcs46HY9bxcMw6no5+zNp88iYAAPAuz/RYAAAA+wgWAAAgYQgWAAAgYQgWAAAgYTwTLGbNmqV+/fopIyNDY8aM0YoVK2yXlBJKSko0atQoZWdnKz8/X9dcc402btwY0+bQoUMqLi5Wt27d1LlzZ33jG9/Qrl27Ytps375dkydPVlZWlvLz83Xvvfeqrq4ups3bb7+tESNGKBgMauDAgXr66aeT/fE8b+bMmXIcR3fffbe7jePV/uzYsUPf+ta31K1bN2VmZmro0KFatWqV+7wxRg888IB69uypzMxMTZgwQZs3b455j/3792vKlCnKyclRXl6evvOd76iqqiqmzQcffKALLrhAGRkZKioq0q9+9as2+XxeU19fr/vvv1/9+/dXZmamTj31VP385z+PuceGp4+Z8YC5c+ea9PR084c//MF8+OGH5nvf+57Jy8szu3btsl2a502cONHMmTPHrF+/3qxZs8ZcccUVpk+fPqaqqsptc9ttt5mioiKzaNEis2rVKvOVr3zFnHfeee7zdXV1ZsiQIWbChAnmn//8p5k3b57p3r27mTFjhttm69atJisry0ybNs1s2LDBPPbYY8bv95s333yzTT+vl6xYscL069fPDBs2zNx1113udo5X+7J//37Tt29fc/PNN5vly5ebrVu3mvnz55stW7a4bWbOnGlyc3PNK6+8YtauXWuuuuoq079/f1NdXe22+epXv2rOPvtss2zZMvP3v//dDBw40Nx4443u8xUVFaagoMBMmTLFrF+/3jz//PMmMzPTPPHEE236eb3goYceMt26dTOvv/662bZtm3nxxRdN586dza9//Wu3jZePmSeCxejRo01xcbG7Xl9fb3r16mVKSkosVpWadu/ebSSZJUuWGGOMKS8vN2lpaebFF19023z00UdGknnvvfeMMcbMmzfP+Hw+U1ZW5raZPXu2ycnJMeFw2BhjzI9//GMzePDgmH1df/31ZuLEicn+SJ5UWVlpTjvtNLNgwQJz0UUXucGC49X+/OQnPzHnn3/+MZ+PRCKmsLDQPPLII+628vJyEwwGzfPPP2+MMWbDhg1Gklm5cqXb5o033jCO45gdO3YYY4z57W9/a7p06eIew+i+zzjjjER/JM+bPHmyufXWW2O2ff3rXzdTpkwxxnj/mHX4oZCamhqtXr1aEyZMcLf5fD5NmDBB7733nsXKUlNFRYUkqWvXrpKk1atXq7a2Nub4DBo0SH369HGPz3vvvaehQ4eqoKDAbTNx4kSFQiF9+OGHbpvm7xFtwzFumeLiYk2ePPmI75Tj1f68+uqrGjlypK677jrl5+dr+PDheuqpp9znt23bprKyspjvOzc3V2PGjIk5Znl5eRo5cqTbZsKECfL5fFq+fLnb5sILL1R6errbZuLEidq4caO+/PLLZH9MTznvvPO0aNEibdq0SZK0du1avfPOO5o0aZIk7x+zNr8JWaLt3btX9fX1Mf/ISVJBQYE+/vhjS1Wlpkgkorvvvlvjxo3TkCFDJEllZWVKT09XXl5eTNuCggKVlZW5bY52/KLPHa9NKBRSdXW1MjMzk/GRPGnu3Ll6//33tXLlyiOe43i1P1u3btXs2bM1bdo03XfffVq5cqXuvPNOpaena+rUqe53frTvu/nxyM/Pj3k+EAioa9euMW369+9/xHtEn+vSpUtSPp8XTZ8+XaFQSIMGDZLf71d9fb0eeughTZkyRZI8f8w6fLBA+1FcXKz169frnXfesV0KjqG0tFR33XWXFixYoIyMDNvl4CREIhGNHDlSDz/8sCRp+PDhWr9+vX73u99p6tSplqvD0bzwwgt69tln9dxzz2nw4MFas2aN7r77bvXq1SsljlmHHwrp3r27/H7/EbPWd+3apcLCQktVpZ477rhDr7/+ut566y317t3b3V5YWKiamhqVl5fHtG9+fAoLC496/KLPHa9NTk4Of/3GYfXq1dq9e7dGjBihQCCgQCCgJUuW6De/+Y0CgYAKCgo4Xu1Mz549ddZZZ8VsO/PMM7V9+3ZJTd/58f4NLCws1O7du2Oer6ur0/79++M6rjg59957r6ZPn64bbrhBQ4cO1be//W3dc889KikpkeT9Y9bhg0V6errOPfdcLVq0yN0WiUS0aNEijR071mJlqcEYozvuuEMvv/yyFi9efES33Lnnnqu0tLSY47Nx40Zt377dPT5jx47VunXrYv4jWrBggXJyctx/UMeOHRvzHtE2HOP4jB8/XuvWrdOaNWvcx8iRIzVlyhR3mePVvowbN+6IU7g3bdqkvn37SpL69++vwsLCmO87FApp+fLlMcesvLxcq1evdtssXrxYkUhEY8aMcdssXbpUtbW1bpsFCxbojDPOYBgkTgcPHpTPF/vz6vf7FYlEJKXAMbM6dTRB5s6da4LBoHn66afNhg0bzPe//32Tl5cXM2sdyXH77beb3Nxc8/bbb5svvvjCfRw8eNBtc9ttt5k+ffqYxYsXm1WrVpmxY8easWPHus9HT1+8/PLLzZo1a8ybb75pevTocdTTF++9917z0UcfmVmzZnH6YoI0PyvEGI5Xe7NixQoTCATMQw89ZDZv3myeffZZk5WVZZ555hm3zcyZM01eXp75y1/+Yj744ANz9dVXH/XUxeHDh5vly5ebd955x5x22mkxpy6Wl5ebgoIC8+1vf9usX7/ezJ0712RlZVk/dbEjmjp1qjnllFPc001feukl0717d/PjH//YbePlY+aJYGGMMY899pjp06ePSU9PN6NHjzbLli2zXVJKkHTUx5w5c9w21dXV5gc/+IHp0qWLycrKMtdee6354osvYt7n008/NZMmTTKZmZmme/fu5l//9V9NbW1tTJu33nrLnHPOOSY9Pd0MGDAgZh9oucODBcer/XnttdfMkCFDTDAYNIMGDTJPPvlkzPORSMTcf//9pqCgwASDQTN+/HizcePGmDb79u0zN954o+ncubPJyckxt9xyi6msrIxps3btWnP++eebYDBoTjnlFDNz5sykfzYvCoVC5q677jJ9+vQxGRkZZsCAAeanP/1pzGmhXj5m3DYdAAAkTIefYwEAANoPggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEgYggUAAEiY/w9ulBuNxEbxegAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# stage 3: train to 32 bits\n",
    "randkey, trainkey = jax.random.split(randkey)\n",
    "model = train_stage(trainkey, model, optimizer, batch_size,\n",
    "                    iterations=32768, min_val=0, max_val=(1 << 32) - 1)\n",
    "test_fizzbuzz(model, [i for i in range(int(2e2), int(2e2) + 25)])\n",
    "test_fizzbuzz(model, [i for i in range(int(1e9), int(1e9) + 25)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cb69a399",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72c15724a9494c519bbef619216c5567",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48949678 neither wrong: expect neither got 0.64756584\n",
      "48949738 neither wrong: expect neither got 0.63570035\n",
      "48949802 neither wrong: expect neither got 0.23076619\n",
      "48949818 fizz wrong: expect fizz got 0.728774\n",
      "48949918 neither wrong: expect neither got 0.48839158\n",
      "48949922 neither wrong: expect neither got 0.31886557\n",
      "48949931 neither wrong: expect neither got 0.38822156\n",
      "48949938 fizz wrong: expect fizz got 0.7156811\n",
      "48949946 neither wrong: expect neither got 0.2861112\n",
      "48949947 fizz wrong: expect fizz got 0.722468\n",
      "48949950 fizz wrong: expect fizz got 0.523748\n",
      "48949978 neither wrong: expect neither got 0.7377738\n",
      "48949994 neither wrong: expect neither got 0.56062484\n",
      "48950008 neither wrong: expect neither got 0.6009756\n",
      "48950010 fizz wrong: expect fizz got 0.68679327\n",
      "48950158 neither wrong: expect neither got 0.45459938\n",
      "48950186 neither wrong: expect neither got 0.34532258\n",
      "48950202 fizz wrong: expect fizz got 0.6460029\n",
      "48950248 neither wrong: expect neither got 0.602318\n",
      "52235340 buzz wrong: expect buzz got 0.6042997\n"
     ]
    }
   ],
   "source": [
    "test_fizzbuzz(model, jnp.arange(0, 101), quiet=True)\n",
    "for i in trange(0, int(64e6), int(4e6)):\n",
    "    start = np.random.randint(i, i + 1048576 - 1024)\n",
    "    test_fizzbuzz(model, jnp.arange(start, start + 1024), quiet=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f4ddb9bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n",
      "fizz\n",
      "4\n",
      "buzz\n",
      "fizz\n",
      "7\n",
      "8\n",
      "fizz\n",
      "buzz\n",
      "11\n",
      "fizz\n",
      "13\n",
      "14\n",
      "fizzbuzz\n",
      "16\n",
      "17\n",
      "fizz\n",
      "19\n",
      "buzz\n",
      "fizz\n",
      "22\n",
      "23\n",
      "fizz\n",
      "buzz\n",
      "26\n",
      "fizz\n",
      "28\n",
      "29\n",
      "fizzbuzz\n",
      "31\n",
      "32\n",
      "fizz\n",
      "34\n",
      "buzz\n",
      "fizz\n",
      "37\n",
      "38\n",
      "fizz\n",
      "buzz\n",
      "41\n",
      "fizz\n",
      "43\n",
      "44\n",
      "fizzbuzz\n",
      "46\n",
      "47\n",
      "fizz\n",
      "49\n",
      "buzz\n",
      "fizz\n",
      "52\n",
      "53\n",
      "fizz\n",
      "buzz\n",
      "56\n",
      "fizz\n",
      "58\n",
      "59\n",
      "fizzbuzz\n",
      "61\n",
      "62\n",
      "fizz\n",
      "64\n",
      "buzz\n",
      "fizz\n",
      "67\n",
      "68\n",
      "fizz\n",
      "buzz\n",
      "71\n",
      "fizz\n",
      "73\n",
      "74\n",
      "fizzbuzz\n",
      "76\n",
      "77\n",
      "fizz\n",
      "79\n",
      "buzz\n",
      "fizz\n",
      "82\n",
      "83\n",
      "fizz\n",
      "buzz\n",
      "86\n",
      "fizz\n",
      "88\n",
      "89\n",
      "fizzbuzz\n",
      "91\n",
      "92\n",
      "fizz\n",
      "94\n",
      "buzz\n",
      "fizz\n",
      "97\n",
      "98\n",
      "fizz\n",
      "buzz\n"
     ]
    }
   ],
   "source": [
    "infer_threshold = 0.75\n",
    "def fizzbuzz(model, min, max):\n",
    "    nums = jnp.arange(min, max)\n",
    "    inputs = split_nums(nums)\n",
    "    outputs = model.forward(inputs)\n",
    "    for num, (fizz, buzz, neither) in zip(nums, outputs):\n",
    "        if fizz > infer_threshold:\n",
    "            print('fizz', end='')\n",
    "        if buzz > infer_threshold:\n",
    "            print('buzz', end='')\n",
    "        if neither > infer_threshold:\n",
    "            print(num, end='')\n",
    "        print('')\n",
    "\n",
    "fizzbuzz(model, 1, 101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "38423617",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save/load\n",
    "import pickle\n",
    "from uuid import uuid4\n",
    "\n",
    "def save_model(model, name=None):\n",
    "    out_name = name or f'model.{uuid4()}.pickle'\n",
    "    with open(out_name, 'wb') as file:\n",
    "        pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    \n",
    "    return out_name\n",
    "\n",
    "def load_model(name):\n",
    "    with open(name, 'rb') as file:\n",
    "        return pickle.load(file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}