JAX_Demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "JAX Demo.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "whig_9Vg_DKf",
        "4t9Fj_Rs7oea",
        "kHjSPCYY7_Gu",
        "ZNlOBQCJ9xsz",
        "43Ts13he-CwE",
        "y3fiLWTcBI9I"
      ],
      "authorship_tag": "ABX9TyPJ5tQbN4pYh5tN/TmC6Kn6",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/n2cholas/dsc-workshops/blob/master/JAX_Demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5b6PacdNwqhp",
        "colab_type": "text"
      },
      "source": [
        "# JAX Demo\n",
        "\n",
        "\n",
        "\n",
        "Here's a brief overview of some key JAX functions, followed by an end-to-end logistic regression example. \n",
        "\n",
        "Created 2 July 2020 for the University of Waterloo's Data Science Club's Introduction to JAX workshop. Slides available [here](https://docs.google.com/presentation/d/11dPbvZVsHBzSpbA227RMZ7O_Cl-Yc1UTtEAwdKa2tTc/edit?usp=sharing). The recording will be made available on [this channel](https://www.youtube.com/channel/UCknY88pglf2xz_S72WHIDxg)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "whig_9Vg_DKf",
        "colab_type": "text"
      },
      "source": [
        "## Set Up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dx5C9cmkoGBr",
        "colab_type": "text"
      },
      "source": [
        "Run this cell if you are using a TPU runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FlWlcmJh--66",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "db777ee6-e9c9-4368-e523-3264f1096658"
      },
      "source": [
        "# Make sure the Colab Runtime is set to Accelerator: TPU.\n",
        "import requests\n",
        "import os\n",
        "if 'TPU_DRIVER_MODE' not in globals():\n",
        "  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20191206'\n",
        "  resp = requests.post(url)\n",
        "  TPU_DRIVER_MODE = 1\n",
        "\n",
        "# The following is required to use TPU Driver as JAX's backend.\n",
        "from jax.config import config\n",
        "config.FLAGS.jax_xla_backend = \"tpu_driver\"\n",
        "config.FLAGS.jax_backend_target = \"grpc://\" + os.environ['COLAB_TPU_ADDR']\n",
        "print(config.FLAGS.jax_backend_target)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "grpc://10.4.236.234:8470\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PlD_HuuaoL0k",
        "colab_type": "text"
      },
      "source": [
        "This code let's you conveniently load MNIST from torchvision."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RNnrijsSoFSj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import torchvision\n",
        "import numpy as np\n",
        "\n",
        "def MNIST(train=False):\n",
        "    transform = torchvision.transforms.Compose([\n",
        "    torchvision.transforms.ToTensor(),\n",
        "    torchvision.transforms.Normalize((0.1307,), (0.3081,)),\n",
        "        lambda x: x.numpy().squeeze(),\n",
        "    ])\n",
        "    return torchvision.datasets.MNIST(os.getcwd(), train=train, \n",
        "                                     download=True, \n",
        "                                     transform=transform)\n",
        "\n",
        "class DataLoader(torch.utils.data.DataLoader):\n",
        "    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,\n",
        "                 batch_sampler=None, num_workers=0, drop_last=False, timeout=0,\n",
        "                 worker_init_fn=None):\n",
        "        super().__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=collate_fn,\n",
        "                         pin_memory=False,  # pinning doesn't affect jax arrays (right?)\n",
        "                         drop_last=drop_last,\n",
        "                         timeout=timeout,\n",
        "                         worker_init_fn=worker_init_fn)\n",
        "\n",
        "\n",
        "def collate_fn(batch):\n",
        "    if isinstance(batch[0], jnp.ndarray):\n",
        "        return np.stack(batch)\n",
        "    elif isinstance(batch[0], (tuple, list)):\n",
        "        return type(batch[0])(collate_fn(samples) for samples in zip(*batch))\n",
        "    elif isinstance(batch[0], dict):\n",
        "        return {k: collate_fn([d[k] for d in batch]) for k in batch[0]}\n",
        "    else:\n",
        "        return np.asarray(batch)"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZMUYwK_47jog",
        "colab_type": "text"
      },
      "source": [
        "# Mini-Demo\n",
        "\n",
        "In this section, we look at how JAX's function transformations work."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6WJDoaM8s4eV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import jax\n",
        "from jax import grad, jit, vmap, pmap, value_and_grad\n",
        "import jax.numpy as jnp\n",
        "\n",
        "def cool_fn(a, b, x):\n",
        "    return a*jnp.dot(b, x)\n",
        "\n",
        "a = jnp.array(3.)\n",
        "b = jnp.array([3., 2., 1.])\n",
        "x = jnp.array([4., 4., 4.])"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTG3O1sFs7F4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e3a72a01-fcc1-4248-c5f3-2ffb6d74b35b"
      },
      "source": [
        "cool_fn(a, b, x)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(72., dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4t9Fj_Rs7oea",
        "colab_type": "text"
      },
      "source": [
        "## grad\n",
        "\n",
        "grad computes the gradient with respect to one or more of the input arguments. By default, it computes the gradient with respect to the first argument."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D41cujTc1pD4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "b6a9538e-3fb7-4fda-d23a-e8b2138a77ee"
      },
      "source": [
        "grad(cool_fn)(a, b, x)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(24., dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUPZrd0opP71",
        "colab_type": "text"
      },
      "source": [
        "We can also compute the gradient with respect to multiple arguments, such as `b` and `x` below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fgrda_y17J2P",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "477e4579-4c60-435c-b8e5-ace493c035bb"
      },
      "source": [
        "grad(cool_fn, argnums=(1, 2))(a, b, x)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray([12., 12., 12.], dtype=float32),\n",
              " DeviceArray([9., 6., 3.], dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jgQkgkF_ppOx",
        "colab_type": "text"
      },
      "source": [
        "We can take a take a higher order derivate by composing the `grad` functions:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AbHEmtra7bHB",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "ed9dc142-47a6-45d0-abf6-221d6e37360e"
      },
      "source": [
        "grad(grad(grad(grad(cool_fn))))(a, b, x)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0., dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RNctrh4rpq9e",
        "colab_type": "text"
      },
      "source": [
        "It's often useful to get both the value and the gradient:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qh_tujS17qFw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "78d8afc3-a444-4d01-9e09-603a5c5bc916"
      },
      "source": [
        "value_and_grad(cool_fn)(a, b, x)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(72., dtype=float32), DeviceArray(24., dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kHjSPCYY7_Gu",
        "colab_type": "text"
      },
      "source": [
        "## vmap\n",
        "\n",
        "Vectorized map allows you to add a batch dimension to one or more arguments:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7l3-80j88m1R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "x2 = jnp.array([[1., 1., 1.],\n",
        "               [2., 3., 1.],\n",
        "               [3., 4., 4.],\n",
        "               [4., 5., 2.]])"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rcr6SWkfo1FV",
        "colab_type": "text"
      },
      "source": [
        "`in_axes` is a tuple representing which axis will be batched for each argument. Here, we indicate only the third argument will be batched, and it will be batched at the 0th dimension."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_xy-Advn8Av6",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "d3f91a93-dcda-4673-f46c-f2ccef062075"
      },
      "source": [
        "vmap(cool_fn, in_axes=(None, None, 0))(a, b, x2)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([18., 39., 63., 72.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xXAA5LDIpvmB",
        "colab_type": "text"
      },
      "source": [
        "We can arbitrarily compute `grad` and `vmap`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oDFWFlnv9c77",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "07529f15-bd48-43ac-871a-67d65c1dc9f3"
      },
      "source": [
        "vmap(grad(cool_fn), in_axes=(None, None, 0))(a, b, x2)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 6., 13., 21., 24.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZNlOBQCJ9xsz",
        "colab_type": "text"
      },
      "source": [
        "## jit\n",
        "\n",
        "`jit` (just in time compilation) compiles your python code via XLA to optimize it and remove python overhead:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DMOTceG29y09",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "93f56d17-9832-46fe-cc4b-2a96d92c4006"
      },
      "source": [
        "jit(cool_fn)(a, b, x)"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(72., dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "veTZKfBup7Kv",
        "colab_type": "text"
      },
      "source": [
        "As with all the function transformations, we can compute them:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7VONv55w90sm",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "e728aebd-b8dc-4e66-9d03-5d681c38c567"
      },
      "source": [
        "jit(vmap(grad(cool_fn), in_axes=(None, None, 0)))(a, b, x2)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 6., 13., 21., 24.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43Ts13he-CwE",
        "colab_type": "text"
      },
      "source": [
        "## pmap\n",
        "\n",
        "Parallel map let's you execute a function on a homogenous set of devices (e.g. multiple TPUs or multiple GPUs)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_8FaMOBd-DWE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "75a4a228-0cf1-4aeb-c254-d0d2465bfeb4"
      },
      "source": [
        "jax.local_devices()"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=2, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=4, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=5, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=6, host_id=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=7, host_id=0, coords=(0,0,0), core_on_chip=0)]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xxrswAkWA339",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "1c10527b-8ba9-4259-ccfb-2956a4ef5ca9"
      },
      "source": [
        "a = jnp.ones(8)\n",
        "a"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VE87XgIbqOB-",
        "colab_type": "text"
      },
      "source": [
        "The parallelized dimension must have elements less than or equal to the number of devices. The input will be split along the parallelized dimension(s), the computation will happen on each device, then the results are concatenated:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0uZLG1vqAcdI",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "3c4eeab6-f597-42c6-a22b-2fe15072cf71"
      },
      "source": [
        "pmap(cool_fn, in_axes=(0, None, None))(a, b, x)"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "ShardedDeviceArray([24., 24., 24., 24., 24., 24., 24., 24.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y3fiLWTcBI9I",
        "colab_type": "text"
      },
      "source": [
        "## randomness\n",
        "\n",
        "In JAX, you must use an explict PRNG key. Check out https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html for more details"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6nbq9fE2qw-Y",
        "colab_type": "text"
      },
      "source": [
        "In numpy, you're used to generating random numbers via an implicit random key:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wbO52FsSBJ03",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kvDh_IIXDhj9",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "6ad6429a-d618-40dd-ed31-6a00fa12bd7e"
      },
      "source": [
        "np.random.rand(5)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([0.18557401, 0.11987591, 0.27819761, 0.43926921, 0.05016992])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UbpdgNsXq0gi",
        "colab_type": "text"
      },
      "source": [
        "In JAX, that's explicit:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0A3y6VZjDjF7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "prng = jax.random.PRNGKey(0)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kowfbHSxq2q5",
        "colab_type": "text"
      },
      "source": [
        "You generally want to use the key once, so I typically split the key before using it, so I have a current key and a key that I will use in the future:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QulYzc_iExCr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "prng, cur = jax.random.split(prng, 2)"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bhba1JVREz4M",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "d742167e-ec31-46a1-94e7-7f104a238d26"
      },
      "source": [
        "jax.random.uniform(cur, (5,))\n",
        "# use prng for other things"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([0.07239354, 0.02032685, 0.9796876 , 0.87867916, 0.16457272],            dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJgGfcgJJqRO",
        "colab_type": "text"
      },
      "source": [
        "# Logistic Regression in JAX\n",
        "\n",
        "We'll write a simple training loop to train a logistic regression model with minibatch gradient descent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Q6EwTGqJs1L",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P05hvk-2rX3A",
        "colab_type": "text"
      },
      "source": [
        "We load MNIST, a dataset of 60000 handwritten digits, using `torchvision`. We just need to augment our dataloader to return numpy arrays instead of torch tensors (in the set up code at the top)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eTTdHTlkrKJ4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dataset = MNIST(train=True)\n",
        "dataloader = DataLoader(dataset, batch_size=128, num_workers=4)"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wIX-t2QRrL7h",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 282
        },
        "outputId": "3e2c443f-cdb2-4321-98a0-3ca8197f535e"
      },
      "source": [
        "image, label = dataset[0]\n",
        "plt.imshow(image)\n",
        "print(label)"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "5\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAOZ0lEQVR4nO3dbYxc5XnG8euKbezamMQbB9chLjjgFAg0Jl0ZEBZQobgOqgSoCsSKIkJpnSY4Ca0rQWlV3IpWbpUQUUqRTHExFS+BBIQ/0CTUQpCowWWhBgwEDMY0NmaNWYENIX5Z3/2w42iBnWeXmTMv3vv/k1Yzc+45c24NXD5nznNmHkeEAIx/H+p0AwDag7ADSRB2IAnCDiRB2IEkJrZzY4d5ckzRtHZuEkjlV3pbe2OPR6o1FXbbiyVdJ2mCpH+LiJWl50/RNJ3qc5rZJICC9bGubq3hw3jbEyTdIOnzkk6UtMT2iY2+HoDWauYz+wJJL0TE5ojYK+lOSedV0xaAqjUT9qMk/WLY4621Ze9ie6ntPtt9+7Snic0BaEbLz8ZHxKqI6I2I3kma3OrNAaijmbBvkzRn2ONP1JYB6ELNhP1RSfNsz7V9mKQvSlpbTVsAqtbw0FtE7Le9TNKPNDT0tjoinq6sMwCVamqcPSLul3R/Rb0AaCEulwWSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiCJpmZxRffzxPJ/4gkfm9nS7T/3F8fUrQ1OPVBc9+hjdxTrU7/uYv3Vaw+rW3u893vFdXcOvl2sn3r38mL9uD9/pFjvhKbCbnuLpN2SBiXtj4jeKpoCUL0q9uy/FxE7K3gdAC3EZ3YgiWbDHpJ+bPsx20tHeoLtpbb7bPft054mNwegUc0exi+MiG22j5T0gO2fR8TDw58QEaskrZKkI9wTTW4PQIOa2rNHxLba7Q5J90paUEVTAKrXcNhtT7M9/eB9SYskbayqMQDVauYwfpake20ffJ3bI+KHlXQ1zkw4YV6xHpMnFeuvnPWRYv2d0+qPCfd8uDxe/JPPlMebO+k/fzm9WP/Hf1lcrK8/+fa6tZf2vVNcd2X/54r1j//k0PtE2nDYI2KzpM9U2AuAFmLoDUiCsANJEHYgCcIOJEHYgST4imsFBs/+bLF+7S03FOufmlT/q5jj2b4YLNb/5vqvFOsT3y4Pf51+97K6tenb9hfXnbyzPDQ3tW99sd6N2LMDSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKMs1dg8nOvFOuP/WpOsf6pSf1VtlOp5dtPK9Y3v1X+Kepbjv1+3dqbB8rj5LP++b+L9VY69L7AOjr27EAShB1IgrADSRB2IAnCDiRB2IEkCDuQhCPaN6J4hHviVJ/Ttu11i4FLTi/Wdy0u/9zzhCcPL9af+Pr1H7ing67Z+TvF+qNnlcfRB994s1iP0+v/APGWbxZX1dwlT5SfgPdZH+u0KwZGnMuaPTuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJME4exeYMPOjxfrg6wPF+ku31x8rf/rM1cV1F/zDN4r1I2/o3HfK8cE1Nc5ue7XtHbY3DlvWY/sB25tqtzOqbBhA9cZyGH+LpPfOen+lpHURMU/SutpjAF1s1LBHxMOS3nsceZ6kNbX7aySdX3FfACrW6G/QzYqI7bX7r0qaVe+JtpdKWipJUzS1wc0BaFbTZ+Nj6Axf3bN8EbEqInojoneSJje7OQANajTs/bZnS1Ltdkd1LQFohUbDvlbSxbX7F0u6r5p2ALTKqJ/Zbd8h6WxJM21vlXS1pJWS7rJ9qaSXJV3YyibHu8Gdrze1/r5djc/v/ukvPVOsv3bjhPILHCjPsY7uMWrYI2JJnRJXxwCHEC6XBZIg7EAShB1IgrADSRB2IAmmbB4HTrji+bq1S04uD5r8+9HrivWzvnBZsT79e48U6+ge7NmBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnG2ceB0rTJr3/thOK6/7f2nWL9ymtuLdb/8sILivX43w/Xrc35+58V11Ubf+Y8A/bsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEUzYnN/BHpxfrt1397WJ97sQpDW/707cuK9bn3bS9WN+/eUvD2x6vmpqyGcD4QNiBJAg7kARhB5Ig7EAShB1IgrADSTDOjqI4Y36xfsTKrcX6HZ/8UcPbPv7BPy7Wf/tv63+PX5IGN21ueNuHqqbG2W2vtr3D9sZhy1bY3mZ7Q+3v3CobBlC9sRzG3yJp8QjLvxsR82t/91fbFoCqjRr2iHhY0kAbegHQQs2coFtm+8naYf6Mek+yvdR2n+2+fdrTxOYANKPRsN8o6VhJ8yVtl/Sdek+MiFUR0RsRvZM0ucHNAWhWQ2GPiP6IGIyIA5JukrSg2rYAVK2hsNuePezhBZI21nsugO4w6ji77TsknS1ppqR+SVfXHs+XFJK2SPpqRJS/fCzG2cejCbOOLNZfuei4urX1V1xXXPdDo+yLvvTSomL9zYWvF+vjUWmcfdRJIiJiyQiLb266KwBtxeWyQBKEHUiCsANJEHYgCcIOJMFXXNExd20tT9k81YcV67+MvcX6H3zj8vqvfe/64rqHKn5KGgBhB7Ig7EAShB1IgrADSRB2IAnCDiQx6rfekNuBheWfkn7xC+Upm0+av6VubbRx9NFcP3BKsT71vr6mXn+8Yc8OJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kwzj7OufekYv35b5bHum86Y02xfuaU8nfKm7En9hXrjwzMLb/AgVF/3TwV9uxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kATj7IeAiXOPLtZfvOTjdWsrLrqzuO4fHr6zoZ6qcFV/b7H+0HWnFesz1pR/dx7vNuqe3fYc2w/afsb207a/VVveY/sB25tqtzNa3y6ARo3lMH6/pOURcaKk0yRdZvtESVdKWhcR8yStqz0G0KVGDXtEbI+Ix2v3d0t6VtJRks6TdPBayjWSzm9VkwCa94E+s9s+RtIpktZLmhURBy8+flXSrDrrLJW0VJKmaGqjfQJo0pjPxts+XNIPJF0eEbuG12JodsgRZ4iMiFUR0RsRvZM0ualmATRuTGG3PUlDQb8tIu6pLe63PbtWny1pR2taBFCFUQ/jbVvSzZKejYhrh5XWSrpY0sra7X0t6XAcmHjMbxXrb/7u7GL9or/7YbH+px+5p1hvpeXby8NjP/vX+sNrPbf8T3HdGQcYWqvSWD6znyHpy5Kesr2htuwqDYX8LtuXSnpZ0oWtaRFAFUYNe0T8VNKIk7tLOqfadgC0CpfLAkkQdiAJwg4kQdiBJAg7kARfcR2jibN/s25tYPW04rpfm/tQsb5ken9DPVVh2baFxfrjN5anbJ75/Y3Fes9uxsq7BXt2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUgizTj73t8v/2zx3j8bKNavOu7+urVFv/F2Qz1VpX/wnbq1M9cuL657/F//vFjveaM8Tn6gWEU3Yc8OJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0mkGWffcn7537XnT767Zdu+4Y1ji/XrHlpUrHuw3o/7Djn+mpfq1ub1ry+uO1isYjxhzw4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTgiyk+w50i6VdIsSSFpVURcZ3uFpD+R9FrtqVdFRP0vfUs6wj1xqpn4FWiV9bFOu2JgxAszxnJRzX5JyyPicdvTJT1m+4Fa7bsR8e2qGgXQOmOZn327pO21+7ttPyvpqFY3BqBaH+gzu+1jJJ0i6eA1mMtsP2l7te0ZddZZarvPdt8+7WmqWQCNG3PYbR8u6QeSLo+IXZJulHSspPka2vN/Z6T1ImJVRPRGRO8kTa6gZQCNGFPYbU/SUNBvi4h7JCki+iNiMCIOSLpJ0oLWtQmgWaOG3bYl3Szp2Yi4dtjy2cOedoGk8nSeADpqLGfjz5D0ZUlP2d5QW3aVpCW252toOG6LpK+2pEMAlRjL2fifShpp3K44pg6gu3AFHZAEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IIlRf0q60o3Zr0l6ediimZJ2tq2BD6Zbe+vWviR6a1SVvR0dER8bqdDWsL9v43ZfRPR2rIGCbu2tW/uS6K1R7eqNw3ggCcIOJNHpsK/q8PZLurW3bu1LordGtaW3jn5mB9A+nd6zA2gTwg4k0ZGw215s+znbL9i+shM91GN7i+2nbG+w3dfhXlbb3mF747BlPbYfsL2pdjviHHsd6m2F7W21926D7XM71Nsc2w/afsb207a/VVve0feu0Fdb3re2f2a3PUHS85I+J2mrpEclLYmIZ9raSB22t0jqjYiOX4Bh+0xJb0m6NSJOqi37J0kDEbGy9g/ljIi4okt6WyHprU5P412brWj28GnGJZ0v6Svq4HtX6OtCteF968SefYGkFyJic0TslXSnpPM60EfXi4iHJQ28Z/F5ktbU7q/R0P8sbVent64QEdsj4vHa/d2SDk4z3tH3rtBXW3Qi7EdJ+sWwx1vVXfO9h6Qf237M9tJONzOCWRGxvXb/VUmzOtnMCEadxrud3jPNeNe8d41Mf94sTtC938KI+Kykz0u6rHa42pVi6DNYN42djmka73YZYZrxX+vke9fo9OfN6kTYt0maM+zxJ2rLukJEbKvd7pB0r7pvKur+gzPo1m53dLifX+umabxHmmZcXfDedXL6806E/VFJ82zPtX2YpC9KWtuBPt7H9rTaiRPZniZpkbpvKuq1ki6u3b9Y0n0d7OVdumUa73rTjKvD713Hpz+PiLb/STpXQ2fkX5T0V53ooU5fn5T0RO3v6U73JukODR3W7dPQuY1LJX1U0jpJmyT9l6SeLurtPyQ9JelJDQVrdod6W6ihQ/QnJW2o/Z3b6feu0Fdb3jculwWS4AQdkARhB5Ig7EAShB1IgrADSRB2IAnCDiTx/65XcTNOWsh5AAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kZ6_HrCrrnpQ",
        "colab_type": "text"
      },
      "source": [
        "We first create our parameters, our weight matrix `W` and our bias `b`. We only use our RNG once, so we don't split it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fbJ04G0HrqgY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "rng = jax.random.PRNGKey(0)\n",
        "W = jax.random.uniform(rng, (28*28, 10))\n",
        "b = jnp.zeros(10)\n",
        "params = (W, b)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wBaMir0arv4n",
        "colab_type": "text"
      },
      "source": [
        "In order to make a prediction, we define an `apply` function. We use a `log_softmax` instead of a softmax for numerical stability."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m64P1RR1r3pF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def apply(params, image):\n",
        "    W, b = params\n",
        "    return jax.nn.log_softmax(W.T @ image.flatten() + b)"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K8HnUAbhr4bZ",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "73a2f1a6-8732-48a6-f61e-617f7f140d9c"
      },
      "source": [
        "jnp.exp(apply(params, image))  # represents the probability of each class"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([2.15396110e-08, 2.32523686e-04, 7.38948316e-08,\n",
              "             1.87424382e-11, 3.16452235e-02, 9.64966595e-01,\n",
              "             5.38649192e-06, 4.69302108e-09, 3.14939930e-03,\n",
              "             1.00649586e-07], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "080cb-aNsHRl",
        "colab_type": "text"
      },
      "source": [
        "We want to feed in batches of images to our apply function, so let's use vmap. The same parameters are used for all the input images and our batch dimension is at the front, hence `in_axes=(None, 0)`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q4deR4rZsQYD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "batch_apply = jax.vmap(apply, in_axes=(None, 0))"
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0VuStMpMsbT1",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e6016bb9-18da-420a-faf8-226f53f88581"
      },
      "source": [
        "images, labels = next(iter(dataloader))\n",
        "logits = batch_apply(params, images)\n",
        "logits.shape"
      ],
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(128, 10)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mPLV-Rdlshlh",
        "colab_type": "text"
      },
      "source": [
        "As we expect, we get our `log_softmax` vector of size 10 for all 128 elements.\n",
        "\n",
        "Next, we define our cross entropy loss function. Below is the definition for a single sample:\n",
        "\n",
        "$$ \\mathcal L(\\hat y, y) = - \\sum_{i=1}^{10} y_i \\log \\hat y_i $$\n",
        "\n",
        "We write it for the entire batch, since it's simple enough without vmap. We return the average over the batch:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EQtuk9M4sfDd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def cross_entropy_loss(logits, labels):\n",
        "    return -(jax.nn.one_hot(labels, 10) * logits).sum(-1).mean()"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lRPFiU-4tE4v",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "1af17e32-5f21-4480-eff0-42a625591217"
      },
      "source": [
        "cross_entropy_loss(logits, labels)"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(11.61635, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msAE8W17tJOn",
        "colab_type": "text"
      },
      "source": [
        "Now, we define the entire loss, which, given parameters and a batch of data, will return the average loss for that batch and the logits (which will be useful to compute accuracy):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J1uHoyQAtcec",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def loss_fn(params, batch):\n",
        "    images, labels = batch\n",
        "    logits = batch_apply(params, images)\n",
        "    return cross_entropy_loss(logits, labels), logits"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gtm9g_SPtcw-",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "14e70a50-e23e-4bc2-e077-4619c696601c"
      },
      "source": [
        "loss, logits = loss_fn(params, (images, labels))\n",
        "print(f'Loss: {loss:.4f}, Logits Shape: {logits.shape}')"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loss: 11.6164, Logits Shape: (128, 10)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rfe1rHn7tr6_",
        "colab_type": "text"
      },
      "source": [
        "We also define accuracy, which is our metric of interest:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IXQaEyy3txD0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def accuracy(logits, labels):\n",
        "    return (logits.argmax(-1) == labels).mean()"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "10tKLiB6tx69",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "673bf307-c2cc-4e3d-ce66-80ba746f0646"
      },
      "source": [
        "accuracy(logits, labels)"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0.1171875, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wLrhIObCt0-s",
        "colab_type": "text"
      },
      "source": [
        "When we fit our parameters to our data, we use stochastic gradient descent (SGD). It's expressed as follows for a single parameter (e.g. `W`):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LcjXbAstuC7o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def sgd_step(param, grad, lr=0.01):\n",
        "    return param - lr * grad"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mryAlUdGuEip",
        "colab_type": "text"
      },
      "source": [
        "We know have all the tools we need to write a single train step:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_oeV3c2lWEDA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "@jax.jit  # compile the top level function, not each individual function\n",
        "def train_step(params, batch):\n",
        "    # has_aux=True is required because we return our loss (what we're differentiating)\n",
        "    # w.r.t, as well as our logits\n",
        "    (loss, logits), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, batch)\n",
        "    # In JAX, any nested structure of dicts, tuples, and lists containing jax\n",
        "    # arrays is called a PyTree. Multimap applies the function to all the arrays\n",
        "    # and returns a PyTree of the same structure.\n",
        "    new_params = jax.tree_multimap(sgd_step, params, grads)\n",
        "    acc = accuracy(logits, batch[1])\n",
        "    return new_params, (loss, acc)"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0EPo4UtIugEf",
        "colab_type": "text"
      },
      "source": [
        "We can now train our model (we'll just train for 300 steps for now)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "svKykWETWdf6",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 544
        },
        "outputId": "79b61c82-3180-4631-d3c8-f11341358b6a"
      },
      "source": [
        "for i, batch in enumerate(dataloader):\n",
        "    params, (loss, acc) = train_step(params, batch)\n",
        "    if i % 10 == 0:\n",
        "        print(f'Iter: {i:3d}  Loss: {loss:.4f}, Accuracy: {acc:.4f}')\n",
        "    if i > 300:\n",
        "        break"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Iter:   0  Loss: 11.6164, Accuracy: 0.1172\n",
            "Iter:  10  Loss: 8.5104, Accuracy: 0.0859\n",
            "Iter:  20  Loss: 6.2342, Accuracy: 0.2266\n",
            "Iter:  30  Loss: 5.3381, Accuracy: 0.2422\n",
            "Iter:  40  Loss: 4.6509, Accuracy: 0.2891\n",
            "Iter:  50  Loss: 4.2726, Accuracy: 0.4062\n",
            "Iter:  60  Loss: 3.9557, Accuracy: 0.3984\n",
            "Iter:  70  Loss: 2.9911, Accuracy: 0.4531\n",
            "Iter:  80  Loss: 3.2631, Accuracy: 0.4688\n",
            "Iter:  90  Loss: 3.0934, Accuracy: 0.4141\n",
            "Iter: 100  Loss: 2.2155, Accuracy: 0.6172\n",
            "Iter: 110  Loss: 2.6738, Accuracy: 0.5312\n",
            "Iter: 120  Loss: 2.3034, Accuracy: 0.5469\n",
            "Iter: 130  Loss: 2.6763, Accuracy: 0.5625\n",
            "Iter: 140  Loss: 2.0505, Accuracy: 0.6250\n",
            "Iter: 150  Loss: 1.9659, Accuracy: 0.6562\n",
            "Iter: 160  Loss: 2.3841, Accuracy: 0.5469\n",
            "Iter: 170  Loss: 1.3181, Accuracy: 0.7188\n",
            "Iter: 180  Loss: 1.6405, Accuracy: 0.6250\n",
            "Iter: 190  Loss: 1.8172, Accuracy: 0.6875\n",
            "Iter: 200  Loss: 1.9610, Accuracy: 0.6484\n",
            "Iter: 210  Loss: 2.1923, Accuracy: 0.5859\n",
            "Iter: 220  Loss: 1.4158, Accuracy: 0.6875\n",
            "Iter: 230  Loss: 1.3246, Accuracy: 0.7266\n",
            "Iter: 240  Loss: 1.9980, Accuracy: 0.6797\n",
            "Iter: 250  Loss: 1.9020, Accuracy: 0.6484\n",
            "Iter: 260  Loss: 1.5684, Accuracy: 0.6172\n",
            "Iter: 270  Loss: 1.7863, Accuracy: 0.6484\n",
            "Iter: 280  Loss: 1.6871, Accuracy: 0.7031\n",
            "Iter: 290  Loss: 1.4895, Accuracy: 0.7344\n",
            "Iter: 300  Loss: 1.5678, Accuracy: 0.6875\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SSnPFj7HvkiN",
        "colab_type": "text"
      },
      "source": [
        "Nice! Our loss is going down and accuracy going up. Let's see how well our (partially) trained model performs an example:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K_poOyE8vkKT",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 282
        },
        "outputId": "1d82faa3-0245-4b4e-cd7f-2f4a4a55f60b"
      },
      "source": [
        "image, label = dataset[0]\n",
        "\n",
        "pred = apply(params, image).argmax()\n",
        "plt.imshow(image)\n",
        "print(f'Prediction: {pred}  Label: {label}')"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Prediction: 5  Label: 5\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAOZ0lEQVR4nO3dbYxc5XnG8euKbezamMQbB9chLjjgFAg0Jl0ZEBZQobgOqgSoCsSKIkJpnSY4Ca0rQWlV3IpWbpUQUUqRTHExFS+BBIQ/0CTUQpCowWWhBgwEDMY0NmaNWYENIX5Z3/2w42iBnWeXmTMv3vv/k1Yzc+45c24NXD5nznNmHkeEAIx/H+p0AwDag7ADSRB2IAnCDiRB2IEkJrZzY4d5ckzRtHZuEkjlV3pbe2OPR6o1FXbbiyVdJ2mCpH+LiJWl50/RNJ3qc5rZJICC9bGubq3hw3jbEyTdIOnzkk6UtMT2iY2+HoDWauYz+wJJL0TE5ojYK+lOSedV0xaAqjUT9qMk/WLY4621Ze9ie6ntPtt9+7Snic0BaEbLz8ZHxKqI6I2I3kma3OrNAaijmbBvkzRn2ONP1JYB6ELNhP1RSfNsz7V9mKQvSlpbTVsAqtbw0FtE7Le9TNKPNDT0tjoinq6sMwCVamqcPSLul3R/Rb0AaCEulwWSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiCJpmZxRffzxPJ/4gkfm9nS7T/3F8fUrQ1OPVBc9+hjdxTrU7/uYv3Vaw+rW3u893vFdXcOvl2sn3r38mL9uD9/pFjvhKbCbnuLpN2SBiXtj4jeKpoCUL0q9uy/FxE7K3gdAC3EZ3YgiWbDHpJ+bPsx20tHeoLtpbb7bPft054mNwegUc0exi+MiG22j5T0gO2fR8TDw58QEaskrZKkI9wTTW4PQIOa2rNHxLba7Q5J90paUEVTAKrXcNhtT7M9/eB9SYskbayqMQDVauYwfpake20ffJ3bI+KHlXQ1zkw4YV6xHpMnFeuvnPWRYv2d0+qPCfd8uDxe/JPPlMebO+k/fzm9WP/Hf1lcrK8/+fa6tZf2vVNcd2X/54r1j//k0PtE2nDYI2KzpM9U2AuAFmLoDUiCsANJEHYgCcIOJEHYgST4imsFBs/+bLF+7S03FOufmlT/q5jj2b4YLNb/5vqvFOsT3y4Pf51+97K6tenb9hfXnbyzPDQ3tW99sd6N2LMDSRB2IAnCDiRB2IEkCDuQBGEHkiDsQBKMs1dg8nOvFOuP/WpOsf6pSf1VtlOp5dtPK9Y3v1X+Kepbjv1+3dqbB8rj5LP++b+L9VY69L7AOjr27EAShB1IgrADSRB2IAnCDiRB2IEkCDuQhCPaN6J4hHviVJ/Ttu11i4FLTi/Wdy0u/9zzhCcPL9af+Pr1H7ing67Z+TvF+qNnlcfRB994s1iP0+v/APGWbxZX1dwlT5SfgPdZH+u0KwZGnMuaPTuQBGEHkiDsQBKEHUiCsANJEHYgCcIOJME4exeYMPOjxfrg6wPF+ku31x8rf/rM1cV1F/zDN4r1I2/o3HfK8cE1Nc5ue7XtHbY3DlvWY/sB25tqtzOqbBhA9cZyGH+LpPfOen+lpHURMU/SutpjAF1s1LBHxMOS3nsceZ6kNbX7aySdX3FfACrW6G/QzYqI7bX7r0qaVe+JtpdKWipJUzS1wc0BaFbTZ+Nj6Axf3bN8EbEqInojoneSJje7OQANajTs/bZnS1Ltdkd1LQFohUbDvlbSxbX7F0u6r5p2ALTKqJ/Zbd8h6WxJM21vlXS1pJWS7rJ9qaSXJV3YyibHu8Gdrze1/r5djc/v/ukvPVOsv3bjhPILHCjPsY7uMWrYI2JJnRJXxwCHEC6XBZIg7EAShB1IgrADSRB2IAmmbB4HTrji+bq1S04uD5r8+9HrivWzvnBZsT79e48U6+ge7NmBJAg7kARhB5Ig7EAShB1IgrADSRB2IAnG2ceB0rTJr3/thOK6/7f2nWL9ymtuLdb/8sILivX43w/Xrc35+58V11Ubf+Y8A/bsQBKEHUiCsANJEHYgCcIOJEHYgSQIO5AEUzYnN/BHpxfrt1397WJ97sQpDW/707cuK9bn3bS9WN+/eUvD2x6vmpqyGcD4QNiBJAg7kARhB5Ig7EAShB1IgrADSTDOjqI4Y36xfsTKrcX6HZ/8UcPbPv7BPy7Wf/tv63+PX5IGN21ueNuHqqbG2W2vtr3D9sZhy1bY3mZ7Q+3v3CobBlC9sRzG3yJp8QjLvxsR82t/91fbFoCqjRr2iHhY0kAbegHQQs2coFtm+8naYf6Mek+yvdR2n+2+fdrTxOYANKPRsN8o6VhJ8yVtl/Sdek+MiFUR0RsRvZM0ucHNAWhWQ2GPiP6IGIyIA5JukrSg2rYAVK2hsNuePezhBZI21nsugO4w6ji77TsknS1ppqR+SVfXHs+XFJK2SPpqRJS/fCzG2cejCbOOLNZfuei4urX1V1xXXPdDo+yLvvTSomL9zYWvF+vjUWmcfdRJIiJiyQiLb266KwBtxeWyQBKEHUiCsANJEHYgCcIOJMFXXNExd20tT9k81YcV67+MvcX6H3zj8vqvfe/64rqHKn5KGgBhB7Ig7EAShB1IgrADSRB2IAnCDiQx6rfekNuBheWfkn7xC+Upm0+av6VubbRx9NFcP3BKsT71vr6mXn+8Yc8OJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kwzj7OufekYv35b5bHum86Y02xfuaU8nfKm7En9hXrjwzMLb/AgVF/3TwV9uxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kATj7IeAiXOPLtZfvOTjdWsrLrqzuO4fHr6zoZ6qcFV/b7H+0HWnFesz1pR/dx7vNuqe3fYc2w/afsb207a/VVveY/sB25tqtzNa3y6ARo3lMH6/pOURcaKk0yRdZvtESVdKWhcR8yStqz0G0KVGDXtEbI+Ix2v3d0t6VtJRks6TdPBayjWSzm9VkwCa94E+s9s+RtIpktZLmhURBy8+flXSrDrrLJW0VJKmaGqjfQJo0pjPxts+XNIPJF0eEbuG12JodsgRZ4iMiFUR0RsRvZM0ualmATRuTGG3PUlDQb8tIu6pLe63PbtWny1pR2taBFCFUQ/jbVvSzZKejYhrh5XWSrpY0sra7X0t6XAcmHjMbxXrb/7u7GL9or/7YbH+px+5p1hvpeXby8NjP/vX+sNrPbf8T3HdGQcYWqvSWD6znyHpy5Kesr2htuwqDYX8LtuXSnpZ0oWtaRFAFUYNe0T8VNKIk7tLOqfadgC0CpfLAkkQdiAJwg4kQdiBJAg7kARfcR2jibN/s25tYPW04rpfm/tQsb5ken9DPVVh2baFxfrjN5anbJ75/Y3Fes9uxsq7BXt2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUgizTj73t8v/2zx3j8bKNavOu7+urVFv/F2Qz1VpX/wnbq1M9cuL657/F//vFjveaM8Tn6gWEU3Yc8OJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0mkGWffcn7537XnT767Zdu+4Y1ji/XrHlpUrHuw3o/7Djn+mpfq1ub1ry+uO1isYjxhzw4kQdiBJAg7kARhB5Ig7EAShB1IgrADSTgiyk+w50i6VdIsSSFpVURcZ3uFpD+R9FrtqVdFRP0vfUs6wj1xqpn4FWiV9bFOu2JgxAszxnJRzX5JyyPicdvTJT1m+4Fa7bsR8e2qGgXQOmOZn327pO21+7ttPyvpqFY3BqBaH+gzu+1jJJ0i6eA1mMtsP2l7te0ZddZZarvPdt8+7WmqWQCNG3PYbR8u6QeSLo+IXZJulHSspPka2vN/Z6T1ImJVRPRGRO8kTa6gZQCNGFPYbU/SUNBvi4h7JCki+iNiMCIOSLpJ0oLWtQmgWaOG3bYl3Szp2Yi4dtjy2cOedoGk8nSeADpqLGfjz5D0ZUlP2d5QW3aVpCW252toOG6LpK+2pEMAlRjL2fifShpp3K44pg6gu3AFHZAEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSRB2IIlRf0q60o3Zr0l6ediimZJ2tq2BD6Zbe+vWviR6a1SVvR0dER8bqdDWsL9v43ZfRPR2rIGCbu2tW/uS6K1R7eqNw3ggCcIOJNHpsK/q8PZLurW3bu1LordGtaW3jn5mB9A+nd6zA2gTwg4k0ZGw215s+znbL9i+shM91GN7i+2nbG+w3dfhXlbb3mF747BlPbYfsL2pdjviHHsd6m2F7W21926D7XM71Nsc2w/afsb207a/VVve0feu0Fdb3re2f2a3PUHS85I+J2mrpEclLYmIZ9raSB22t0jqjYiOX4Bh+0xJb0m6NSJOqi37J0kDEbGy9g/ljIi4okt6WyHprU5P412brWj28GnGJZ0v6Svq4HtX6OtCteF968SefYGkFyJic0TslXSnpPM60EfXi4iHJQ28Z/F5ktbU7q/R0P8sbVent64QEdsj4vHa/d2SDk4z3tH3rtBXW3Qi7EdJ+sWwx1vVXfO9h6Qf237M9tJONzOCWRGxvXb/VUmzOtnMCEadxrud3jPNeNe8d41Mf94sTtC938KI+Kykz0u6rHa42pVi6DNYN42djmka73YZYZrxX+vke9fo9OfN6kTYt0maM+zxJ2rLukJEbKvd7pB0r7pvKur+gzPo1m53dLifX+umabxHmmZcXfDedXL6806E/VFJ82zPtX2YpC9KWtuBPt7H9rTaiRPZniZpkbpvKuq1ki6u3b9Y0n0d7OVdumUa73rTjKvD713Hpz+PiLb/STpXQ2fkX5T0V53ooU5fn5T0RO3v6U73JukODR3W7dPQuY1LJX1U0jpJmyT9l6SeLurtPyQ9JelJDQVrdod6W6ihQ/QnJW2o/Z3b6feu0Fdb3jculwWS4AQdkARhB5Ig7EAShB1IgrADSRB2IAnCDiTx/65XcTNOWsh5AAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W6jDKKv_uq1D",
        "colab_type": "text"
      },
      "source": [
        "You can extend this example by making use of all the TPU cores instead of just one by incorporating `pmap`. There's a guide here: https://github.com/google/jax/blob/master/cloud_tpu_colabs/Pmap_Cookbook.ipynb. \n",
        "\n",
        "Check out https://github.com/google/jax/ for more complete examples. "
      ]
    }
  ]
}