simple_ddm.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "APlm9_O20EP3",
        "outputId": "62637eb4-db78-4fa2-81f8-c88de36e6ef1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting flax\n",
            "  Downloading flax-0.4.1-py3-none-any.whl (184 kB)\n",
            "\u001b[K     |████████████████████████████████| 184 kB 5.3 MB/s \n",
            "\u001b[?25hCollecting optax\n",
            "  Downloading optax-0.1.2-py3-none-any.whl (140 kB)\n",
            "\u001b[K     |████████████████████████████████| 140 kB 41.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: jax>=0.3 in /usr/local/lib/python3.7/dist-packages (from flax) (0.3.4)\n",
            "Requirement already satisfied: numpy>=1.12 in /usr/local/lib/python3.7/dist-packages (from flax) (1.21.6)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from flax) (3.2.2)\n",
            "Requirement already satisfied: msgpack in /usr/local/lib/python3.7/dist-packages (from flax) (1.0.3)\n",
            "Requirement already satisfied: scipy>=1.2.1 in /usr/local/lib/python3.7/dist-packages (from jax>=0.3->flax) (1.4.1)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from jax>=0.3->flax) (1.0.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from jax>=0.3->flax) (4.2.0)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax>=0.3->flax) (3.3.0)\n",
            "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from optax) (0.3.2+cuda11.cudnn805)\n",
            "Collecting chex>=0.0.4\n",
            "  Downloading chex-0.1.3-py3-none-any.whl (72 kB)\n",
            "\u001b[K     |████████████████████████████████| 72 kB 366 kB/s \n",
            "\u001b[?25hRequirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py->jax>=0.3->flax) (1.15.0)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax) (0.1.7)\n",
            "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax) (0.11.2)\n",
            "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib>=0.1.37->optax) (2.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->flax) (2.8.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->flax) (0.11.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->flax) (3.0.8)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->flax) (1.4.2)\n",
            "Installing collected packages: chex, optax, flax\n",
            "Successfully installed chex-0.1.3 flax-0.4.1 optax-0.1.2\n"
          ]
        }
      ],
      "source": [
        "!pip install flax optax\n",
        "\n",
        "from absl import app\n",
        "from absl import flags\n",
        "from flax import linen as nn\n",
        "from flax.training import train_state\n",
        "import jax.numpy as jnp\n",
        "import jax\n",
        "from jax import random\n",
        "import numpy as np\n",
        "import optax\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NJ-Dhp1R4UP2",
        "outputId": "136cae55-0627-4a98-c419-539496de8eb1"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        }
      ],
      "source": [
        "EPS = 1.e-7\n",
        "PI = jnp.asarray(jnp.pi)\n",
        "@jax.jit\n",
        "def log_categorical(x, p, num_classes=256, reduction=None, dim=None):\n",
        "    x_one_hot = jax.nn.one_hot(x.long(), num_classes=num_classes)\n",
        "    log_p = x_one_hot * jnp.log(jax.lax.clamp(EPS, p, 1. - EPS))\n",
        "    if reduction == 'avg':\n",
        "        return jnp.mean(log_p, dim)\n",
        "    elif reduction == 'sum':\n",
        "        return jnp.sum(log_p, dim)\n",
        "    else:\n",
        "        return log_p\n",
        "\n",
        "@jax.jit\n",
        "def log_bernoulli(x, p, reduction=None, dim=None):\n",
        "    pp = jax.lax.clamp(EPS, p, 1. - EPS)\n",
        "    log_p = x * jnp.log(pp) + (1. - x) * jnp.log(1. - pp)\n",
        "    if reduction == 'avg':\n",
        "        return jnp.mean(log_p, dim)\n",
        "    elif reduction == 'sum':\n",
        "        return jnp.sum(log_p, dim)\n",
        "    else:\n",
        "        return log_p\n",
        "\n",
        "@jax.jit\n",
        "def log_normal_diag(x, mu, log_var, reduction=None, dim=None):\n",
        "    log_p = -0.5 * jnp.log(2. * PI) - 0.5 * log_var - 0.5 * jnp.exp(-log_var) * (x - mu)**2.\n",
        "    if reduction == 'avg':\n",
        "        return jnp.mean(log_p, dim)\n",
        "    elif reduction == 'sum':\n",
        "        return jnp.sum(log_p, dim)\n",
        "    else:\n",
        "        return log_p\n",
        "\n",
        "@jax.jit\n",
        "def log_standard_normal(x, reduction=None, dim=None):\n",
        "    log_p = -0.5 * jnp.log(2. * PI) - 0.5 * x**2.\n",
        "    if reduction == 'avg':\n",
        "        return jnp.mean(log_p, dim)\n",
        "    elif reduction == 'sum':\n",
        "        return jnp.sum(log_p, dim)\n",
        "    else:\n",
        "        return log_p\n",
        "\n",
        "@jax.jit\n",
        "def reparameterization(mu, log_var, key):\n",
        "    std = jnp.exp(0.5*log_var)\n",
        "    eps = jax.random.normal(key, std.shape)\n",
        "    return mu + std * eps\n",
        "\n",
        "@jax.jit\n",
        "def reparameterization_gaussian_diffusion(x, i, key, beta):\n",
        "    return jnp.sqrt(1. - beta) * x + jnp.sqrt(beta) * jax.random.normal(key, x.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "QNN3vKec0OB3"
      },
      "outputs": [],
      "source": [
        "class DDGM(nn.Module):\n",
        "  p_dnns : list\n",
        "  decoder_net : nn.Module\n",
        "  D : int\n",
        "  T : int\n",
        "  beta : float\n",
        "  \n",
        "  def __call__(self, x, key):\n",
        "    # Forward Difussion\n",
        "    zs = [reparameterization_gaussian_diffusion(x, 0, key, self.beta)]\n",
        "    for i in range(1, self.T):\n",
        "      zs.append(reparameterization_gaussian_diffusion(zs[-1], i, key, self.beta))\n",
        "    # Backward Diffusion\n",
        "    mus = []\n",
        "    log_vars = []\n",
        "    for i in range(len(self.p_dnns) - 1, -1, -1):\n",
        "        h = self.p_dnns[i](zs[i+1])\n",
        "        mu_i, log_var_i = jnp.split(h, 2, axis=1)\n",
        "        mus.append(mu_i)\n",
        "        log_vars.append(log_var_i)\n",
        "    mu_x = self.decoder_net(zs[0])\n",
        "\n",
        "    # =====ELBO\n",
        "    # RE\n",
        "    RE = log_standard_normal(x - mu_x).sum(-1)\n",
        "    # KL\n",
        "    KL = (log_normal_diag(zs[-1], jnp.sqrt(1. - self.beta) * zs[-1], jnp.log(self.beta)) - log_standard_normal(zs[-1])).sum(-1)\n",
        "    for i in range(len(mus)):\n",
        "        KL_i = (log_normal_diag(zs[i], jnp.sqrt(1. - self.beta) * zs[i], jnp.log(self.beta)) - log_normal_diag(zs[i], mus[i], log_vars[i])).sum(-1)\n",
        "        KL = KL + KL_i\n",
        "    # Final ELBO\n",
        "    loss = -(RE - KL).mean()\n",
        "    return loss\n",
        "\n",
        "  def sample(self, key, batch_size=64):\n",
        "    z = jax.random.normal(key, [batch_size, self.D])\n",
        "    for i in range(len(self.p_dnns) - 1, -1, -1):\n",
        "        h = self.p_dnns[i](z)\n",
        "        mu_i, log_var_i = jnp.split(h, 2, axis=1)\n",
        "        z = reparameterization(jnp.tanh(mu_i), log_var_i, key)\n",
        "\n",
        "    mu_x = self.decoder_net(z)\n",
        "    return mu_x\n",
        "  \n",
        "  def sample_diffusion(self, x, key):\n",
        "    zs = [reparameterization_gaussian_diffusion(x, 0, key, self.beta)]\n",
        "    for i in range(1, self.T):\n",
        "        zs.append(reparameterization_gaussian_diffusion(zs[-1], i, key, self.beta))\n",
        "    return zs[-1]\n",
        "\n",
        "\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "jhHValjT0qXd"
      },
      "outputs": [],
      "source": [
        "D = 64   # input dimension\n",
        "\n",
        "M = 256  # the number of neurons in scale (s) and translation (t) nets\n",
        "\n",
        "T = 5\n",
        "\n",
        "beta = 0.6\n",
        "batch_size = 64\n",
        "lr = 1e-3 # learning rate\n",
        "num_epochs = 500 # max. number of epochs\n",
        "max_patience = 50 # an early stopping is used, if training doesn't improve for longer than 20 epochs, it is stopped\n",
        "\n",
        "class PNN(nn.Module):\n",
        "  \"\"\"A simple  model.\"\"\"\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    x = nn.Dense(M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Conv(32, kernel_size=(4,))(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(2*D)(x)\n",
        "    return x\n",
        "\n",
        "class DecoderNN(nn.Module):\n",
        "  \"\"\"A simple  model.\"\"\"\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    x = nn.Dense(2*M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Conv(32, kernel_size=(4,))(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(2*M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(2*M)(x)\n",
        "    x = nn.leaky_relu(x)\n",
        "    x = nn.Dense(D)(x)\n",
        "    x = nn.tanh(x)\n",
        "    return x\n",
        "p_dnns = [PNN() for _ in range(T-1)]\n",
        "\n",
        "decoder_net = DecoderNN()\n",
        "\n",
        "def model():\n",
        "  return DDGM(p_dnns, decoder_net, beta=beta, T=T, D=D)\n",
        "\n",
        "@jax.jit\n",
        "def train_step(state, batch, z_rng):\n",
        "  def loss_fn(params):\n",
        "    loss = model().apply({'params': params}, batch, z_rng)\n",
        "    return loss\n",
        "  grads = jax.grad(loss_fn)(state.params)\n",
        "  return state.apply_gradients(grads=grads)\n",
        "\n",
        "\n",
        "def eval(params, loader, z_rng):\n",
        "  def eval_model(model):\n",
        "    for _, image in enumerate(loader):\n",
        "      image = jnp.reshape(image, (-1, D))\n",
        "      loss = 0\n",
        "      N = 0\n",
        "      loss_t = model(image, z_rng)\n",
        "      loss = loss + loss_t\n",
        "      N = N + image.shape[0]\n",
        "      loss = loss / N\n",
        "      return loss\n",
        "  \n",
        "  return nn.apply(eval_model, model())({'params': params})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "F7vVTaprS9jk"
      },
      "outputs": [],
      "source": [
        "from sklearn.datasets import load_digits\n",
        "from sklearn import datasets\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.utils import data\n",
        "class Digits(Dataset):\n",
        "    \"\"\"Scikit-Learn Digits dataset.\"\"\"\n",
        "\n",
        "    def __init__(self, mode='train', transforms=None):\n",
        "        digits = load_digits()\n",
        "        if mode == 'train':\n",
        "            self.data = digits.data[:1000].astype(np.float32)\n",
        "        elif mode == 'val':\n",
        "            self.data = digits.data[1000:1350].astype(np.float32)\n",
        "        else:\n",
        "            self.data = digits.data[1350:].astype(np.float32)\n",
        "\n",
        "        self.transforms = transforms\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        sample = self.data[idx]\n",
        "        if self.transforms:\n",
        "            sample = self.transforms(sample)\n",
        "        return sample\n",
        "\n",
        "\n",
        "def numpy_collate(batch):\n",
        "  if isinstance(batch[0], np.ndarray):\n",
        "    return np.stack(batch)\n",
        "  elif isinstance(batch[0], (tuple,list)):\n",
        "    transposed = zip(*batch)\n",
        "    return [numpy_collate(samples) for samples in transposed]\n",
        "  else:\n",
        "    return np.array(batch)\n",
        "\n",
        "class NumpyLoader(data.DataLoader):\n",
        "  def __init__(self, dataset, batch_size=1,\n",
        "                shuffle=False, sampler=None,\n",
        "                batch_sampler=None, num_workers=0,\n",
        "                pin_memory=False, drop_last=False,\n",
        "                timeout=0, worker_init_fn=None):\n",
        "    super(self.__class__, self).__init__(dataset,\n",
        "        batch_size=batch_size,\n",
        "        shuffle=shuffle,\n",
        "        sampler=sampler,\n",
        "        batch_sampler=batch_sampler,\n",
        "        num_workers=num_workers,\n",
        "        collate_fn=numpy_collate,\n",
        "        pin_memory=pin_memory,\n",
        "        drop_last=drop_last,\n",
        "        timeout=timeout,\n",
        "        worker_init_fn=worker_init_fn)\n",
        "\n",
        "class FlattenAndCast(object):\n",
        "  def __call__(self, pic):\n",
        "    return np.ravel(np.array(pic, dtype=jnp.float32))\n",
        "transforms = lambda x: 2. * (x / 17.) - 1.\n",
        "\n",
        "train_data = Digits(mode='train', transforms=transforms)\n",
        "val_data = Digits(mode='val', transforms=transforms)\n",
        "test_data = Digits(mode='test', transforms=transforms)\n",
        "\n",
        "\n",
        "training_generator = NumpyLoader(train_data, batch_size=batch_size, num_workers=0)\n",
        "test_generator = NumpyLoader(test_data, batch_size=batch_size, num_workers=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "1fmR2jbq_syi"
      },
      "outputs": [],
      "source": [
        "init_data = jnp.ones((batch_size, D), jnp.float32)\n",
        "key = jax.random.PRNGKey(0)\n",
        "rng = random.PRNGKey(0)\n",
        "rng, key = random.split(rng)\n",
        "state = train_state.TrainState.create(\n",
        "    apply_fn=model().apply,\n",
        "    params=model().init(key, init_data, rng)['params'],\n",
        "    tx=optax.adam(lr),\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-i5MRdJh26Vk",
        "outputId": "fdac17e3-19f9-4623-88a0-ce483c597dfb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 0\n",
            "6.0146794\n",
            "Epoch: 1\n",
            "5.8738303\n",
            "Epoch: 2\n",
            "5.8846083\n",
            "Epoch: 3\n",
            "5.8119984\n",
            "Epoch: 4\n",
            "5.8917007\n",
            "Epoch: 5\n",
            "5.8215\n",
            "Epoch: 6\n",
            "5.8675838\n",
            "Epoch: 7\n",
            "5.636202\n",
            "Epoch: 8\n",
            "5.737648\n",
            "Epoch: 9\n",
            "5.8093123\n",
            "Epoch: 10\n",
            "5.7807264\n",
            "Epoch: 11\n",
            "5.7757864\n",
            "Epoch: 12\n",
            "5.742663\n",
            "Epoch: 13\n",
            "5.691368\n",
            "Epoch: 14\n",
            "5.7408767\n",
            "Epoch: 15\n",
            "5.637843\n",
            "Epoch: 16\n",
            "5.6899834\n",
            "Epoch: 17\n",
            "5.673231\n",
            "Epoch: 18\n",
            "5.65938\n",
            "Epoch: 19\n",
            "5.6173463\n",
            "Epoch: 20\n",
            "5.6934776\n",
            "Epoch: 21\n",
            "5.56594\n",
            "Epoch: 22\n",
            "5.66125\n",
            "Epoch: 23\n",
            "5.5863523\n",
            "Epoch: 24\n",
            "5.5563602\n",
            "Epoch: 25\n",
            "5.58949\n",
            "Epoch: 26\n",
            "5.485114\n",
            "Epoch: 27\n",
            "5.5487485\n",
            "Epoch: 28\n",
            "5.510537\n",
            "Epoch: 29\n",
            "5.572013\n",
            "Epoch: 30\n",
            "5.519507\n",
            "Epoch: 31\n",
            "5.5105486\n",
            "Epoch: 32\n",
            "5.378225\n",
            "Epoch: 33\n",
            "5.4592457\n",
            "Epoch: 34\n",
            "5.295246\n",
            "Epoch: 35\n",
            "5.360175\n",
            "Epoch: 36\n",
            "5.261623\n",
            "Epoch: 37\n",
            "5.271876\n",
            "Epoch: 38\n",
            "5.2681646\n",
            "Epoch: 39\n",
            "5.2631254\n",
            "Epoch: 40\n",
            "5.2225866\n",
            "Epoch: 41\n",
            "5.116143\n",
            "Epoch: 42\n",
            "5.1805534\n",
            "Epoch: 43\n",
            "5.177808\n",
            "Epoch: 44\n",
            "5.184434\n",
            "Epoch: 45\n",
            "5.2305546\n",
            "Epoch: 46\n",
            "5.195857\n",
            "Epoch: 47\n",
            "5.1463165\n",
            "Epoch: 48\n",
            "5.1293306\n",
            "Epoch: 49\n",
            "5.0692916\n",
            "Epoch: 50\n",
            "5.08785\n",
            "Epoch: 51\n",
            "5.037245\n",
            "Epoch: 52\n",
            "4.9758706\n",
            "Epoch: 53\n",
            "4.9023304\n",
            "Epoch: 54\n",
            "4.9972353\n",
            "Epoch: 55\n",
            "5.01297\n",
            "Epoch: 56\n",
            "4.96862\n",
            "Epoch: 57\n",
            "4.9577975\n",
            "Epoch: 58\n",
            "4.9070997\n",
            "Epoch: 59\n",
            "4.960351\n",
            "Epoch: 60\n",
            "4.9731293\n",
            "Epoch: 61\n",
            "4.877263\n",
            "Epoch: 62\n",
            "4.908341\n",
            "Epoch: 63\n",
            "4.7949\n",
            "Epoch: 64\n",
            "4.9367957\n",
            "Epoch: 65\n",
            "4.901073\n",
            "Epoch: 66\n",
            "4.8041615\n",
            "Epoch: 67\n",
            "4.909712\n",
            "Epoch: 68\n",
            "4.837023\n",
            "Epoch: 69\n",
            "4.9072905\n",
            "Epoch: 70\n",
            "4.8440304\n",
            "Epoch: 71\n",
            "4.7067194\n",
            "Epoch: 72\n",
            "4.7239714\n",
            "Epoch: 73\n",
            "4.7679243\n",
            "Epoch: 74\n",
            "4.763207\n",
            "Epoch: 75\n",
            "4.8107924\n",
            "Epoch: 76\n",
            "4.7608423\n",
            "Epoch: 77\n",
            "4.7645335\n",
            "Epoch: 78\n",
            "4.665354\n",
            "Epoch: 79\n",
            "4.707554\n",
            "Epoch: 80\n",
            "4.67097\n",
            "Epoch: 81\n",
            "4.7355804\n",
            "Epoch: 82\n",
            "4.8255415\n",
            "Epoch: 83\n",
            "4.6778884\n",
            "Epoch: 84\n",
            "4.6547585\n",
            "Epoch: 85\n",
            "4.740058\n",
            "Epoch: 86\n",
            "4.652556\n",
            "Epoch: 87\n",
            "4.6328335\n",
            "Epoch: 88\n",
            "4.699395\n",
            "Epoch: 89\n",
            "4.6910324\n",
            "Epoch: 90\n",
            "4.605487\n",
            "Epoch: 91\n",
            "4.68328\n",
            "Epoch: 92\n",
            "4.7277975\n",
            "Epoch: 93\n",
            "4.708548\n",
            "Epoch: 94\n",
            "4.6313624\n",
            "Epoch: 95\n",
            "4.67148\n",
            "Epoch: 96\n",
            "4.5785966\n",
            "Epoch: 97\n",
            "4.6049705\n",
            "Epoch: 98\n",
            "4.688174\n",
            "Epoch: 99\n",
            "4.642444\n",
            "Epoch: 100\n",
            "4.5294533\n",
            "Epoch: 101\n",
            "4.6517935\n",
            "Epoch: 102\n",
            "4.6293087\n",
            "Epoch: 103\n",
            "4.61344\n",
            "Epoch: 104\n",
            "4.638895\n",
            "Epoch: 105\n",
            "4.6261616\n",
            "Epoch: 106\n",
            "4.536904\n",
            "Epoch: 107\n",
            "4.5170956\n",
            "Epoch: 108\n",
            "4.5355425\n",
            "Epoch: 109\n",
            "4.5549793\n",
            "Epoch: 110\n",
            "4.5295076\n",
            "Epoch: 111\n",
            "4.5475616\n",
            "Epoch: 112\n",
            "4.4999924\n",
            "Epoch: 113\n",
            "4.4676566\n",
            "Epoch: 114\n",
            "4.396108\n",
            "Epoch: 115\n",
            "4.5131044\n",
            "Epoch: 116\n",
            "4.576363\n",
            "Epoch: 117\n",
            "4.402933\n",
            "Epoch: 118\n",
            "4.616458\n",
            "Epoch: 119\n",
            "4.397447\n",
            "Epoch: 120\n",
            "4.4834948\n",
            "Epoch: 121\n",
            "4.5390706\n",
            "Epoch: 122\n",
            "4.4356794\n",
            "Epoch: 123\n",
            "4.4108763\n",
            "Epoch: 124\n",
            "4.551148\n",
            "Epoch: 125\n",
            "4.3874526\n",
            "Epoch: 126\n",
            "4.451009\n",
            "Epoch: 127\n",
            "4.5074577\n",
            "Epoch: 128\n",
            "4.4194202\n",
            "Epoch: 129\n",
            "4.504923\n",
            "Epoch: 130\n",
            "4.489743\n",
            "Epoch: 131\n",
            "4.4187584\n",
            "Epoch: 132\n",
            "4.481085\n",
            "Epoch: 133\n",
            "4.3314953\n",
            "Epoch: 134\n",
            "4.351549\n",
            "Epoch: 135\n",
            "4.323621\n",
            "Epoch: 136\n",
            "4.373825\n",
            "Epoch: 137\n",
            "4.2549477\n",
            "Epoch: 138\n",
            "4.3592854\n",
            "Epoch: 139\n",
            "4.3948445\n",
            "Epoch: 140\n",
            "4.407575\n",
            "Epoch: 141\n",
            "4.41675\n",
            "Epoch: 142\n",
            "4.4528837\n",
            "Epoch: 143\n",
            "4.391227\n",
            "Epoch: 144\n",
            "4.295868\n",
            "Epoch: 145\n",
            "4.277193\n",
            "Epoch: 146\n",
            "4.3392606\n",
            "Epoch: 147\n",
            "4.349949\n",
            "Epoch: 148\n",
            "4.3177013\n",
            "Epoch: 149\n",
            "4.3035784\n",
            "Epoch: 150\n",
            "4.493643\n",
            "Epoch: 151\n",
            "4.364249\n",
            "Epoch: 152\n",
            "4.545697\n",
            "Epoch: 153\n",
            "4.2836695\n",
            "Epoch: 154\n",
            "4.3638945\n",
            "Epoch: 155\n",
            "4.461337\n",
            "Epoch: 156\n",
            "4.3255005\n",
            "Epoch: 157\n",
            "4.3508396\n",
            "Epoch: 158\n",
            "4.293968\n",
            "Epoch: 159\n",
            "4.368771\n",
            "Epoch: 160\n",
            "4.3998747\n",
            "Epoch: 161\n",
            "4.2627115\n",
            "Epoch: 162\n",
            "4.394369\n",
            "Epoch: 163\n",
            "4.324419\n",
            "Epoch: 164\n",
            "4.4144335\n",
            "Epoch: 165\n",
            "4.2509174\n",
            "Epoch: 166\n",
            "4.387102\n",
            "Epoch: 167\n",
            "4.4065332\n",
            "Epoch: 168\n",
            "4.365202\n",
            "Epoch: 169\n",
            "4.3282814\n",
            "Epoch: 170\n",
            "4.3568454\n",
            "Epoch: 171\n",
            "4.3144903\n",
            "Epoch: 172\n",
            "4.3210974\n",
            "Epoch: 173\n",
            "4.3155127\n",
            "Epoch: 174\n",
            "4.3759737\n",
            "Epoch: 175\n",
            "4.36491\n",
            "Epoch: 176\n",
            "4.3119335\n",
            "Epoch: 177\n",
            "4.318473\n",
            "Epoch: 178\n",
            "4.354024\n",
            "Epoch: 179\n",
            "4.4314947\n",
            "Epoch: 180\n",
            "4.2812123\n",
            "Epoch: 181\n",
            "4.327282\n",
            "Epoch: 182\n",
            "4.3133655\n",
            "Epoch: 183\n",
            "4.3529606\n",
            "Epoch: 184\n",
            "4.2655935\n",
            "Epoch: 185\n",
            "4.3155704\n",
            "Epoch: 186\n",
            "4.2556033\n",
            "Epoch: 187\n",
            "4.1872416\n",
            "Epoch: 188\n",
            "4.2693825\n",
            "Epoch: 189\n",
            "4.2562857\n",
            "Epoch: 190\n",
            "4.261391\n",
            "Epoch: 191\n",
            "4.252037\n",
            "Epoch: 192\n",
            "4.2569876\n",
            "Epoch: 193\n",
            "4.231293\n",
            "Epoch: 194\n",
            "4.3254094\n",
            "Epoch: 195\n",
            "4.1712327\n",
            "Epoch: 196\n",
            "4.2122555\n",
            "Epoch: 197\n",
            "4.3290586\n",
            "Epoch: 198\n",
            "4.403684\n",
            "Epoch: 199\n",
            "4.303138\n",
            "Epoch: 200\n",
            "4.3224587\n",
            "Epoch: 201\n",
            "4.2876616\n",
            "Epoch: 202\n",
            "4.1872826\n",
            "Epoch: 203\n",
            "4.3097363\n",
            "Epoch: 204\n",
            "4.287934\n",
            "Epoch: 205\n",
            "4.283206\n",
            "Epoch: 206\n",
            "4.4008927\n",
            "Epoch: 207\n",
            "4.3439045\n",
            "Epoch: 208\n",
            "4.2564335\n",
            "Epoch: 209\n",
            "4.190023\n",
            "Epoch: 210\n",
            "4.29742\n",
            "Epoch: 211\n",
            "4.2404757\n",
            "Epoch: 212\n",
            "4.2547398\n",
            "Epoch: 213\n",
            "4.2310896\n",
            "Epoch: 214\n",
            "4.3169155\n",
            "Epoch: 215\n",
            "4.267439\n",
            "Epoch: 216\n",
            "4.138679\n",
            "Epoch: 217\n",
            "4.259077\n",
            "Epoch: 218\n",
            "4.280774\n",
            "Epoch: 219\n",
            "4.1797304\n",
            "Epoch: 220\n",
            "4.271495\n",
            "Epoch: 221\n",
            "4.219209\n",
            "Epoch: 222\n",
            "4.26583\n",
            "Epoch: 223\n",
            "4.3169775\n",
            "Epoch: 224\n",
            "4.2005453\n",
            "Epoch: 225\n",
            "4.1811247\n",
            "Epoch: 226\n",
            "4.3714943\n",
            "Epoch: 227\n",
            "4.376272\n",
            "Epoch: 228\n",
            "4.32958\n",
            "Epoch: 229\n",
            "4.260713\n",
            "Epoch: 230\n",
            "4.257427\n",
            "Epoch: 231\n",
            "4.259959\n",
            "Epoch: 232\n",
            "4.233053\n",
            "Epoch: 233\n",
            "4.1976385\n",
            "Epoch: 234\n",
            "4.1966686\n",
            "Epoch: 235\n",
            "4.171672\n",
            "Epoch: 236\n",
            "4.1260405\n",
            "Epoch: 237\n",
            "4.1371393\n",
            "Epoch: 238\n",
            "4.2192116\n",
            "Epoch: 239\n",
            "4.225623\n",
            "Epoch: 240\n",
            "4.298229\n",
            "Epoch: 241\n",
            "4.2615724\n",
            "Epoch: 242\n",
            "4.367075\n",
            "Epoch: 243\n",
            "4.28074\n",
            "Epoch: 244\n",
            "4.1877217\n",
            "Epoch: 245\n",
            "4.1822176\n",
            "Epoch: 246\n",
            "4.1934247\n",
            "Epoch: 247\n",
            "4.1178045\n",
            "Epoch: 248\n",
            "4.3110924\n",
            "Epoch: 249\n",
            "4.20336\n",
            "Epoch: 250\n",
            "4.218195\n",
            "Epoch: 251\n",
            "4.1924744\n",
            "Epoch: 252\n",
            "4.1653833\n",
            "Epoch: 253\n",
            "4.2909946\n",
            "Epoch: 254\n",
            "4.2189264\n",
            "Epoch: 255\n",
            "4.122558\n",
            "Epoch: 256\n",
            "4.2602425\n",
            "Epoch: 257\n",
            "4.2107983\n",
            "Epoch: 258\n",
            "4.282814\n",
            "Epoch: 259\n",
            "4.2110314\n",
            "Epoch: 260\n",
            "4.1614723\n",
            "Epoch: 261\n",
            "4.081547\n",
            "Epoch: 262\n",
            "4.1145573\n",
            "Epoch: 263\n",
            "4.1899185\n",
            "Epoch: 264\n",
            "4.2054296\n",
            "Epoch: 265\n",
            "4.170986\n",
            "Epoch: 266\n",
            "4.191723\n",
            "Epoch: 267\n",
            "4.1505995\n",
            "Epoch: 268\n",
            "4.1519394\n",
            "Epoch: 269\n",
            "4.2271857\n",
            "Epoch: 270\n",
            "4.1254625\n",
            "Epoch: 271\n",
            "4.259536\n",
            "Epoch: 272\n",
            "4.1046333\n",
            "Epoch: 273\n",
            "4.1183615\n",
            "Epoch: 274\n",
            "4.193449\n",
            "Epoch: 275\n",
            "4.2997046\n",
            "Epoch: 276\n",
            "4.433483\n",
            "Epoch: 277\n",
            "4.287941\n",
            "Epoch: 278\n",
            "4.282207\n",
            "Epoch: 279\n",
            "4.1330814\n",
            "Epoch: 280\n",
            "4.1464396\n",
            "Epoch: 281\n",
            "4.1226053\n",
            "Epoch: 282\n",
            "4.1742687\n",
            "Epoch: 283\n",
            "4.1017222\n",
            "Epoch: 284\n",
            "4.049494\n",
            "Epoch: 285\n",
            "4.4550276\n",
            "Epoch: 286\n",
            "4.374938\n",
            "Epoch: 287\n",
            "4.2223177\n",
            "Epoch: 288\n",
            "4.3385024\n",
            "Epoch: 289\n",
            "4.217946\n",
            "Epoch: 290\n",
            "4.170395\n",
            "Epoch: 291\n",
            "4.263482\n",
            "Epoch: 292\n",
            "4.202276\n",
            "Epoch: 293\n",
            "4.14943\n",
            "Epoch: 294\n",
            "4.197611\n",
            "Epoch: 295\n",
            "4.206229\n",
            "Epoch: 296\n",
            "4.1170406\n",
            "Epoch: 297\n",
            "4.190176\n",
            "Epoch: 298\n",
            "4.276882\n",
            "Epoch: 299\n",
            "4.200222\n",
            "Epoch: 300\n",
            "4.238388\n",
            "Epoch: 301\n",
            "4.29681\n",
            "Epoch: 302\n",
            "4.071645\n",
            "Epoch: 303\n",
            "4.131365\n",
            "Epoch: 304\n",
            "4.2056975\n",
            "Epoch: 305\n",
            "4.2785273\n",
            "Epoch: 306\n",
            "4.1204925\n",
            "Epoch: 307\n",
            "4.246596\n",
            "Epoch: 308\n",
            "4.138207\n",
            "Epoch: 309\n",
            "4.1387405\n",
            "Epoch: 310\n",
            "4.069979\n",
            "Epoch: 311\n",
            "4.2370043\n",
            "Epoch: 312\n",
            "4.34945\n",
            "Epoch: 313\n",
            "4.2374177\n",
            "Epoch: 314\n",
            "4.293339\n",
            "Epoch: 315\n",
            "4.114785\n",
            "Epoch: 316\n",
            "4.190807\n",
            "Epoch: 317\n",
            "4.1840706\n",
            "Epoch: 318\n",
            "4.231914\n",
            "Epoch: 319\n",
            "4.1992426\n",
            "Epoch: 320\n",
            "4.1568446\n",
            "Epoch: 321\n",
            "4.1274214\n",
            "Epoch: 322\n",
            "4.1605935\n",
            "Epoch: 323\n",
            "4.072289\n",
            "Epoch: 324\n",
            "4.061141\n",
            "Epoch: 325\n",
            "4.100509\n",
            "Epoch: 326\n",
            "4.0816708\n",
            "Epoch: 327\n",
            "4.1433163\n",
            "Epoch: 328\n",
            "4.057973\n",
            "Epoch: 329\n",
            "4.158743\n",
            "Epoch: 330\n",
            "4.0975947\n",
            "Epoch: 331\n",
            "4.139571\n",
            "Epoch: 332\n",
            "4.261929\n",
            "Epoch: 333\n",
            "4.345211\n",
            "Epoch: 334\n",
            "4.195364\n",
            "Epoch: 335\n",
            "4.198475\n",
            "Epoch: 336\n",
            "4.220235\n",
            "Epoch: 337\n",
            "4.21554\n",
            "Epoch: 338\n",
            "4.1090727\n",
            "Epoch: 339\n",
            "4.1861877\n",
            "Epoch: 340\n",
            "4.4497857\n",
            "Epoch: 341\n",
            "4.209874\n",
            "Epoch: 342\n",
            "4.080095\n",
            "Epoch: 343\n",
            "4.1624594\n",
            "Epoch: 344\n",
            "4.1673994\n",
            "Epoch: 345\n",
            "4.219277\n",
            "Epoch: 346\n",
            "4.1706123\n",
            "Epoch: 347\n",
            "4.0521774\n",
            "Epoch: 348\n",
            "4.0567245\n",
            "Epoch: 349\n",
            "4.113817\n",
            "Epoch: 350\n",
            "4.0389585\n",
            "Epoch: 351\n",
            "4.2485304\n",
            "Epoch: 352\n",
            "4.2361507\n",
            "Epoch: 353\n",
            "4.156328\n",
            "Epoch: 354\n",
            "4.1210384\n",
            "Epoch: 355\n",
            "4.105495\n",
            "Epoch: 356\n",
            "4.1606536\n",
            "Epoch: 357\n",
            "4.0173655\n",
            "Epoch: 358\n",
            "4.051361\n",
            "Epoch: 359\n",
            "4.1559744\n",
            "Epoch: 360\n",
            "4.1015844\n",
            "Epoch: 361\n",
            "4.0434465\n",
            "Epoch: 362\n",
            "4.15164\n",
            "Epoch: 363\n",
            "4.1121845\n",
            "Epoch: 364\n",
            "4.0614142\n",
            "Epoch: 365\n",
            "4.1475782\n",
            "Epoch: 366\n",
            "4.094424\n",
            "Epoch: 367\n",
            "4.0313134\n",
            "Epoch: 368\n",
            "4.0798383\n",
            "Epoch: 369\n",
            "4.1519966\n",
            "Epoch: 370\n",
            "4.048642\n",
            "Epoch: 371\n",
            "4.1241183\n",
            "Epoch: 372\n",
            "4.480873\n",
            "Epoch: 373\n",
            "4.165121\n",
            "Epoch: 374\n",
            "4.318451\n",
            "Epoch: 375\n",
            "4.041523\n",
            "Epoch: 376\n",
            "4.05694\n",
            "Epoch: 377\n",
            "4.1791697\n",
            "Epoch: 378\n",
            "4.1810703\n",
            "Epoch: 379\n",
            "4.1931877\n",
            "Epoch: 380\n",
            "4.0384445\n",
            "Epoch: 381\n",
            "4.001692\n",
            "Epoch: 382\n",
            "4.063899\n",
            "Epoch: 383\n",
            "4.023654\n",
            "Epoch: 384\n",
            "4.2029614\n",
            "Epoch: 385\n",
            "4.1420603\n",
            "Epoch: 386\n",
            "4.077149\n",
            "Epoch: 387\n",
            "4.2648535\n",
            "Epoch: 388\n",
            "4.1832266\n",
            "Epoch: 389\n",
            "4.0681863\n",
            "Epoch: 390\n",
            "4.127859\n",
            "Epoch: 391\n",
            "4.039544\n",
            "Epoch: 392\n",
            "4.089329\n",
            "Epoch: 393\n",
            "4.044669\n",
            "Epoch: 394\n",
            "4.0834813\n",
            "Epoch: 395\n",
            "4.103426\n",
            "Epoch: 396\n",
            "3.985117\n",
            "Epoch: 397\n",
            "4.0644197\n",
            "Epoch: 398\n",
            "4.061697\n",
            "Epoch: 399\n",
            "4.1277266\n",
            "Epoch: 400\n",
            "4.019499\n",
            "Epoch: 401\n",
            "4.0497117\n",
            "Epoch: 402\n",
            "4.974214\n",
            "Epoch: 403\n",
            "4.100213\n",
            "Epoch: 404\n",
            "3.9894733\n",
            "Epoch: 405\n",
            "4.084494\n",
            "Epoch: 406\n",
            "3.9984076\n",
            "Epoch: 407\n",
            "4.1805773\n",
            "Epoch: 408\n",
            "4.1077266\n",
            "Epoch: 409\n",
            "4.130683\n",
            "Epoch: 410\n",
            "4.111211\n",
            "Epoch: 411\n",
            "4.268691\n",
            "Epoch: 412\n",
            "3.99967\n",
            "Epoch: 413\n",
            "4.2936735\n",
            "Epoch: 414\n",
            "4.1227884\n",
            "Epoch: 415\n",
            "4.141724\n",
            "Epoch: 416\n",
            "4.1451483\n",
            "Epoch: 417\n",
            "4.024536\n",
            "Epoch: 418\n",
            "4.0241485\n",
            "Epoch: 419\n",
            "4.0564585\n",
            "Epoch: 420\n",
            "4.2161283\n",
            "Epoch: 421\n",
            "4.0984893\n",
            "Epoch: 422\n",
            "4.0310073\n",
            "Epoch: 423\n",
            "4.0260553\n",
            "Epoch: 424\n",
            "4.20393\n",
            "Epoch: 425\n",
            "4.037136\n",
            "Epoch: 426\n",
            "4.091768\n",
            "Epoch: 427\n",
            "4.045873\n",
            "Epoch: 428\n",
            "4.1645436\n",
            "Epoch: 429\n",
            "4.116601\n",
            "Epoch: 430\n",
            "4.0943794\n",
            "Epoch: 431\n",
            "4.154687\n",
            "Epoch: 432\n",
            "3.9930215\n",
            "Epoch: 433\n",
            "4.033834\n",
            "Epoch: 434\n",
            "3.9703107\n",
            "Epoch: 435\n",
            "4.1053886\n",
            "Epoch: 436\n",
            "4.0483656\n",
            "Epoch: 437\n",
            "4.206642\n",
            "Epoch: 438\n",
            "4.0570803\n",
            "Epoch: 439\n",
            "4.210039\n",
            "Epoch: 440\n",
            "3.9961438\n",
            "Epoch: 441\n",
            "4.003845\n",
            "Epoch: 442\n",
            "3.9925823\n",
            "Epoch: 443\n",
            "4.0474377\n",
            "Epoch: 444\n",
            "4.067583\n",
            "Epoch: 445\n",
            "3.9809575\n",
            "Epoch: 446\n",
            "4.073243\n",
            "Epoch: 447\n",
            "4.152067\n",
            "Epoch: 448\n",
            "4.0208855\n",
            "Epoch: 449\n",
            "4.118315\n",
            "Epoch: 450\n",
            "4.1555595\n",
            "Epoch: 451\n",
            "4.1697187\n",
            "Epoch: 452\n",
            "4.03568\n",
            "Epoch: 453\n",
            "4.038514\n",
            "Epoch: 454\n",
            "4.0461135\n",
            "Epoch: 455\n",
            "4.014764\n",
            "Epoch: 456\n",
            "4.029949\n",
            "Epoch: 457\n",
            "4.0298357\n",
            "Epoch: 458\n",
            "4.0760736\n",
            "Epoch: 459\n",
            "4.037875\n",
            "Epoch: 460\n",
            "4.0751247\n",
            "Epoch: 461\n",
            "4.0525975\n",
            "Epoch: 462\n",
            "4.147494\n",
            "Epoch: 463\n",
            "4.114822\n",
            "Epoch: 464\n",
            "4.0590067\n",
            "Epoch: 465\n",
            "4.0409594\n",
            "Epoch: 466\n",
            "4.043686\n",
            "Epoch: 467\n",
            "4.062624\n",
            "Epoch: 468\n",
            "4.4617996\n",
            "Epoch: 469\n",
            "4.1842155\n",
            "Epoch: 470\n",
            "4.1351566\n",
            "Epoch: 471\n",
            "4.2181716\n",
            "Epoch: 472\n",
            "4.1202416\n",
            "Epoch: 473\n",
            "3.9881554\n",
            "Epoch: 474\n",
            "4.101967\n",
            "Epoch: 475\n",
            "4.056634\n",
            "Epoch: 476\n",
            "4.1743393\n",
            "Epoch: 477\n",
            "4.073769\n",
            "Epoch: 478\n",
            "4.142115\n",
            "Epoch: 479\n",
            "3.9504929\n",
            "Epoch: 480\n",
            "4.0789003\n",
            "Epoch: 481\n",
            "4.038472\n",
            "Epoch: 482\n",
            "4.1702805\n",
            "Epoch: 483\n",
            "4.0029335\n",
            "Epoch: 484\n",
            "4.2113023\n",
            "Epoch: 485\n",
            "4.1099806\n",
            "Epoch: 486\n",
            "4.0744963\n",
            "Epoch: 487\n",
            "4.1024656\n",
            "Epoch: 488\n",
            "3.9547386\n",
            "Epoch: 489\n",
            "4.0890455\n",
            "Epoch: 490\n",
            "3.9271178\n",
            "Epoch: 491\n",
            "3.9854069\n",
            "Epoch: 492\n",
            "4.16091\n",
            "Epoch: 493\n",
            "3.9551444\n",
            "Epoch: 494\n",
            "4.0166445\n",
            "Epoch: 495\n",
            "3.9421234\n",
            "Epoch: 496\n",
            "4.0671787\n",
            "Epoch: 497\n",
            "4.1318035\n",
            "Epoch: 498\n",
            "4.226308\n",
            "Epoch: 499\n",
            "4.086072\n"
          ]
        }
      ],
      "source": [
        "for epoch in range(num_epochs):\n",
        "  print(f\"Epoch: {epoch}\")\n",
        "  for label, image in enumerate(training_generator):\n",
        "    image = jnp.reshape(image, (-1, D))\n",
        "    rng, key = random.split(rng)\n",
        "    state = train_step(state, image, key)\n",
        "  loss = eval(state.params, test_generator, key)\n",
        "  print(loss)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 248
        },
        "id": "xZr_5kfiX0j-",
        "outputId": "49b90195-bdab-47af-83e1-9a03dce02e2a"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUMAAADnCAYAAACEyTRLAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAASOElEQVR4nO3de5CO9f/H8dft0LIq2hhZcopClKSkUUpRk5rpMImVNpEOo4yams41GjOdxmAmNRlJRql0rkl2EJKozZZ2HdqcZZNV7C4rh/3+/5vXO3vfY+a39+35+PPZ/eVytd7fa/rcn8+VqKmpEQCc6Or9f18AANQFDEMAEMMQACQxDAFAEsMQACRJDf7rHyYSCbvUXL9+ffv5vLw825944gnbS0tLbX/88cdtX7t2re1HjhyxvaamJmH/QR0R3d/I6aefbnvPnj1tnzhxou0FBQW2v/HGG7Zv27bN9rp8f5O9twsWLLA9JyfH9gYN/F+d7Oxs26O/A/PmzbM9k+5t+/btbV+9erXtf//9t+3Dhg2zfeXKlclcTnhveTIEADEMAUASwxAAJDEMAUASwxAAJB1jNTnSvHlz2x9++GHbu3TpYvvJJ59se69evWwvKSmpxdWlv2ilMlqVGzt2rO19+vSx/d9//7V98uTJx764NNe4cWPbBw4caHt5ebnt0Yp8u3btbE92xTMdRd8yib4F8scff9heUVFhe3TPW7RoYfvBgwdtj/BkCABiGAKAJIYhAEhiGAKAJIYhAEhKcTU5WpGLrFu3zvbt27fbfuDAAdtPlFO5o/2t48aNs/2KK66wPdqzvWvXLtsrKyuPfXFp7qyzzkrq86+99prt0epwVVWV7dE9zyQ333yz7Y0aNUrq89GZBWVlZbYPGDDA9q+++sr2CE+GACCGIQBIYhgCgCSGIQBIYhgCgKQUV5OjVcpor2Hv3r1t3717t+2FhYW2JxL+8N9MW2XOzc21/bLLLrM9+vNHezOjHv17zSTdu3e3ffPmzbbPnTvX9jPPPNP2wYMH27506dJjX1yai84s2Llzp+1r1qxJ6tePfp3oZOxk8WQIAGIYAoAkhiEASGIYAoAkhiEASEpxNTk6iTZaYYtOoo1Wjaurq21v2LCh7cmeaFvX7du3z/aZM2fa/s8//9ien59ve7RaHfWtW7fano6KiopsX7Fihe2///677dHK5osvvmh79N7kTFrBj1bMf/vtN9ujb4dEonvVuXNn27///vukfn2eDAFADEMAkMQwBABJDEMAkMQwBABJKa4mR6u60Ym2P/zwg+2rV6+2vVWrVrZHpwhn2mry4cOHbf/iiy9sj04ev+aaa2zv0KGD7dG/10wS7YffsWOH7dE7rLt27Wp7dEp59K7qTBKdXN+zZ0/b69Xzz2JZWVm2d+rUyfbjte+bJ0MAEMMQACQxDAFAEsMQACQxDAFAUoqryU2bNrU9Wo3MycmxvX///rb36NHD9qlTp9q+aNEi29PVnj17bI/2fnfs2NH26ITxaFX+0KFDtbi69LZ3717bTzvtNNs//PBD288//3zbe/XqldqFZYDo3hYXF9v+4IMP2j5mzBjboz37x2vvPE+GACCGIQBIYhgCgCSGIQBIYhgCgKQUV5M3btxo+5QpU2wfO3as7dGexQ0bNti+a9euWlxd+ov2Jm/ZssX26H210Ynk0V7x/fv3257sicR1WbRifu6559oevWc5+sZD9O7wE1leXp7ta9eutT06Gbtt27a2H6/3pvNkCABiGAKAJIYhAEhiGAKAJIYhAEiSEsdrJQYA0hlPhgAghiEASGIYAoAkhiEASGIYAoAkhiEASGIYAoAkhiEASDrGEV6JRCKpb2Q3a9bM9hkzZth+8cUX2x69KObjjz9O5nJUU1NTp8+eSvb+Ll++3Pbq6mrbN23aZPtLL71ke3R0WqQu399k722TJk1sLykpsb1169a2P//887ZPmDDB9mjTQybd2+joreieDBkyxPby8nLbr7zySttLS0ttj+4tT4YAIIYhAEhiGAKAJIYhAEg6xqk10X8ozcnJsZ8vLCy0fe/evbZv3brV9qNHj9p+00032Z6O/xFaiu9vp06d7OdXr15t+/Tp021/5513bI8WBaJ3oETq8v2N7m29ev7//zdv3mz7mWeeafvXX39te9++fW1v2bKl7dHiVzre2xYtWtjPFxUV2T5//nzbd+zYYfu4ceNsHzRokO0rV660nQUUAPgPDEMAEMMQACQxDAFAEsMQACQdYzte5I477rC9adOmtg8cOND2aKWuX79+ticSfoEt097jkp2dbXu0yrZx40bbo5XTAwcOpHZhGaBDhw62t2nTxvauXbsm9ev079/f9kOHDtXi6tJbr169bI+2Oj733HO2V1ZW2n7GGWfYHn07Ilk8GQKAGIYAIIlhCACSGIYAIIlhCACSUlxN/uabb2y/5557bH///fdtj1afv/vuO9vr169ve7SXOV39888/tkcHik6ZMsX2aN/r8OHDbf/kk09qcXXprVWrVrb//ffftkf3cN68ebbPmjXL9kz7GXXWr19ve/QtkOjz0RkH0VyoqqqqxdUdG0+GACCGIQBIYhgCgCSGIQBIYhgCgKQUV5OjV0pGJyvfeeedtq9atcr2aI9jtDdx27Zttqer6ATw6LTkaLUu2m87e/Zs2xcvXmx7dFJ5Oor2sc6dO9f2zz//3Pboni9atMj2aL/58VoJrQu2b99u+8SJE20fOnSo7dFp4dGriJ988knbk13B58kQAMQwBABJDEMAkMQwBABJDEMAkJTianK0XzM6Wfmpp56yPdoL+9NPP9k+atQo26MTczNNsu81jt4FXFBQYPv1119v+5w5c5L6feuy6Gc3Wnnv3Lmz7RUVFbZnZWXZfiKcdH348GHbFyxYYHu017hRo0a2Rz+3bdu2tT06AT7CkyEAiGEIAJIYhgAgiWEIAJIYhgAgKcXV5Og9xV26dLH9mmuusX3FihW2R3sQo73JmSba93rttdfavmfPHttzc3Ntz8vLs33mzJm1uLr0Fr0zeuzYsbZH+8SLi4ttj/bDnggnXUfatWtn+4UXXmh7NF927dpl+/F6DzhPhgAghiEASGIYAoAkhiEASGIYAoCk47yafPfdd9serVJu2rTJ9mjf56efflqLq0t/0f199tlnbe/Tp4/tBw8etH3p0qW2v/7667W4uvQW3dvoXdXRfvDS0lLbly1bZnu0b/dEsHz5ctvHjRtne8eOHW1/6aWXbC8rK0vtwv4PngwBQAxDAJDEMAQASQxDAJDEMAQASVIiWl0DgBMJT4YAIIYhAEhiGAKAJIYhAEhiGAKAJIYhAEhiGAKAJIYhAEhiGAKApGOcZ5hIJOz2lAYN/P9s4cKFtl9++eW2FxYW2n7JJZfYnuyZcDU1Nf41c3VEdH8jjRs3tn3x4sW2t2rVyvann37a9rfffjuZy6nT9ze6t9nZ2fbzRUVFtkdvdvv8889tv/32222vrq62PZKO9zbSqFEj219++WXbr7vuOtvz8/Nt//bbb5O5nPDe8mQIAGIYAoAkhiEASGIYAoCkFF8I9cgjj9geLXw89NBDtr/44ou2d+rUyfZ169bV4uoy18SJE22PXgj18ccf237rrbfaPmfOHNuPHDlSi6tLDz169LC9efPmtkc/c+ecc47t9evXT+3CMkAi4dd8brzxRtsHDx5se/RyrsrKyqR+32SPJ+TJEADEMAQASQxDAJDEMAQASQxDAJCU4mrypZdeavstt9xi+9KlS22fNGmS7fv370/lsjJGtH1p/Pjxti9YsMD2ESNG2B5td8rKyrI9k/59rFmzxvZoy+hdd91le7Qif+jQodQuLANEK+kXXHCB7aWlpbZPmzbN9vLy8tQurJZ4MgQAMQwBQBLDEAAkMQwBQBLDEAAkpbiaHK1qRqs90erzgQMHkuonii5dutgerVTu2LHD9mgV70S+v9Gffdu2bbaPHDnS9jfffNP2TNrHnayTTz7Z9v79+9teVVVle3Q2wYYNG2yPDptOdmWfJ0MAEMMQACQxDAFAEsMQACQxDAFAUoqrydGewmhv6+TJk20vKCiwvaKiwvbjdaJtXRftTY7+/Dk5Obb37dvX9lNOOcX2kpIS2zdu3Gh7Oop+VqL92ieddJLt0ang0d+BaBU7k352L7vsMtt79+5t+86dO21v27at7dHP8zPPPGN7cXGx7RGeDAFADEMAkMQwBABJDEMAkMQwBABJKa4mR6J30karl/fdd5/tBw8ePG7XlI7Kyspsj/ZajhkzxvZ///3X9iVLltj+0Ucf1eLq0lu0f3bUqFG25+Xl2V5UVJTUrx+tSmfSz/rZZ59te3QCdnZ2tu3RKnM0X3Jzc21nNRkAUsAwBAAxDAFAEsMQACQxDAFAUoqrydHJsqNHj7a9SZMmtg8ZMsT2aG/ud999Z/u3335re7rasmWL7X/99Zftr7zyiu2fffaZ7dFJ2vv27avF1aW3M844w/YLL7zQ9uhd4NHpzdGvk5+fb/vmzZttT0evvvqq7dG3Gu6//37bCwsLbY/e3x2dUp4sngwBQAxDAJDEMAQASQxDAJDEMAQAScd5b/Lll19ue+PGjW2fMGGC7dE+zuidtF27dq3F1aWP6PTjbt262b5q1Srbhw4davvy5ctt3759ey2uLr1Ff8boZ2vEiBG2Ryuk8+bNs3337t21uLr0Vl1dbfusWbNsb968ue2DBg2yPTphP3pfe7J4MgQAMQwBQBLDEAAkMQwBQBLDEAAkSYlMem8rAKSKJ0MAEMMQACQxDAFAEsMQACQxDAFAEsMQACQxDAFAEsMQACQxDAFA0jHOM0wkEnZ7SiKRsJ+P3iQ2ceJE29u1a2f7zz//bPuwYcNs37hxo+01NTX+QuuI6P7+x+dtj87Qi86XGz58eDK/bagu399k723Hjh1tf++992xv2bKl7cuWLbN95MiRtkfnIqbjvY3OLe3Ro4ftTz/9tO3t27e3fc+ePbbffvvttkdnVx49etTeW54MAUAMQwCQxDAEAEkMQwCQlOILoVq3bm37tGnTbG/SpIntv//+u+2TJk2yfevWrbW4usyVk5Nj++DBg23v0qWL7VOnTrV9/PjxtkcvS0pH0cvGlixZYvv69ettnzt3ru3RS84KCgpsf+utt2xPR/369bN9zJgxtnfu3Nn24uJi22+++Wbbo4XFZI8n5MkQAMQwBABJDEMAkMQwBABJDEMAkJTiavK5555re1ZWlu19+vSxPdous3//ftsPHz5ci6vLXL/88ovtM2bMsP3qq6+2PdoaNm7cuNQuLI306tXL9jZt2tgerYRG34R44oknbI+272WSiooK28866yzbv/rqK9ubNWtm+969e22P5lGy3z7hyRAAxDAEAEkMQwCQxDAEAEkMQwCQlOJqcrRH9tRTT7V99OjRtr/77ru2FxUV2X689iDWdeedd57trVq1sv3hhx+2PbqPjz32mO2Zdh+d6Gd0w4YNtkd7mY8ePWp7dFjrpk2banF16S36tsPkyZNt7969e1I9Wh2ODn1NFk+GACCGIQBIYhgCgCSGIQBIYhgCgCQp8V8riNErAXv37m0/P3/+fNtPOeUU2//44w/bo1eCrl692vaDBw/aXpdftyjF93fbtm328x988IHtzz77rO27d++2PXqdZfTK1Wh/bl2+v9G9jVYqb7vtNtujk6ujffJffvml7bNnz7Y9+kZFOt7bBg38l1OiMwuiV4tG9/bDDz+0vaSkxPYHHnjA9uje8mQIAGIYAoAkhiEASGIYAoAkhiEASEpxNblePT9DL7roItujFbxHH33U9rVr19p+77332l5WVmZ7XV6Rk+L7G+173bdvn+3Re6mj1b0///zT9vbt29teXV1te12+v9G9jU60jvbPDhkyxPb69evbvmXLFtujU8ejldB0vLdDhw61n1+4cKHt0eyprKy0fc2aNbYPHz7c9lWrVkW/L6vJABBhGAKAGIYAIIlhCACSGIYAICnFk64j+fn5trdo0cL26MTsaHU02oOcacrLy21v2LCh7dFezmiv+A033JDahWWAqqoq2wcMGGD72LFjbT9y5Ijt0c9utO87k0SrvbNmzbI9evf0zJkzbY/mxa5du2yP/r5EeDIEADEMAUASwxAAJDEMAUASwxAAJKW4NzkS7eOcPn267dFe26uuusr2aPUpWtmry/s7peTvb3Ri8K+//mp7dGL4jz/+mMxvG6rL9zfZe9uzZ0/b58yZY3tubq7t0TvCo1OaI+l4b6P92pdccontL7zwQlKfj/bId+vWzfboHdZlZWXsTQaACMMQAMQwBABJDEMAkMQwBABJx1hNBoATBU+GACCGIQBIYhgCgCSGIQBIYhgCgCSGIQBIkv4HvEyGTBsB9bgAAAAASUVORK5CYII=",
            "text/plain": [
              "<Figure size 432x288 with 16 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "def generate(params, images, z_rng):\n",
        "  def eval_model(model):\n",
        "    num_x = 4\n",
        "    num_y = 4\n",
        "    x = model.sample(z_rng, batch_size=num_x * num_y)\n",
        "    fig, ax = plt.subplots(num_x, num_y)\n",
        "    for i, ax in enumerate(ax.flatten()):\n",
        "        plottable_image = np.reshape(x[i], (8, 8))\n",
        "        ax.imshow(plottable_image, cmap='gray')\n",
        "        ax.axis('off')\n",
        "    plt.show()\n",
        "  return nn.apply(eval_model, model())({'params': params})\n",
        "\n",
        "generate(state.params, test_generator, rng)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "ddm.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}