jax_basic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "jax-basic.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPehf0/MS0woXFWJe5XDjUT",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deterministic-algorithms-lab/Jax-Journey/blob/main/jax_basic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oXjPlcOX2G3I"
      },
      "source": [
        "A notebook for this [blog](https://roberttlange.github.io/posts/2020/03/blog-post-10/) with additional notes. Implements MLP and CNN in ```JAX```. It is suggested to read that blog side-by-side. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5Q-brUn68dX"
      },
      "source": [
        "%matplotlib inline\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "import numpy as onp\n",
        "import jax.numpy as np\n",
        "from jax import grad, jit, vmap, value_and_grad\n",
        "from jax import random\n",
        "\n",
        "# Generate key which is used to generate random numbers\n",
        "key = random.PRNGKey(1)                                                         #A key is always an nd-array of size (2,) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U1I15wsT7C_H"
      },
      "source": [
        "# Generate a random matrix\n",
        "x = random.uniform(key, (1000, 1000))\n",
        "# Compare running times of 3 different matrix multiplications\n",
        "%time y = onp.dot(x, x)\n",
        "%time y = np.dot(x, x); print(y)\n",
        "%time y = np.dot(x, x).block_until_ready()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-SbZO63W7_st"
      },
      "source": [
        "The above is due to [Asyncronous dispatch](https://jax.readthedocs.io/en/latest/async_dispatch.html)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i5YUeW4V7QPR"
      },
      "source": [
        "def ReLU(x):\n",
        "    \"\"\" Rectified Linear Unit (ReLU) activation function \"\"\"\n",
        "    return np.maximum(0, x)\n",
        "\n",
        "jit_ReLU = jit(ReLU)                                                            "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uLDbzLKu9HZQ"
      },
      "source": [
        "JIT a simple python function using numpy to make it faster. Normally, each operation has its own kernel which are dispatched to GPU, one by one. If we have a sequence of operations, we can use the ```@jit decorator / jit()``` to compile multiple operations together using XLA."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dm0Oo83u8TdW",
        "outputId": "4025d3d7-6a90-4a54-f951-ff7926162976"
      },
      "source": [
        "%time out = ReLU(x).block_until_ready()\n",
        "\n",
        "# Call jitted version to compile for evaluation time!\n",
        "%time jit_ReLU(x).block_until_ready()                                           #First time call will cause compilation, and may take longer.\n",
        "%time out = jit_ReLU(x).block_until_ready()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 60.6 ms, sys: 0 ns, total: 60.6 ms\n",
            "Wall time: 61.1 ms\n",
            "CPU times: user 26 ms, sys: 841 µs, total: 26.8 ms\n",
            "Wall time: 25.8 ms\n",
            "CPU times: user 1.21 ms, sys: 0 ns, total: 1.21 ms\n",
            "Wall time: 696 µs\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pIMKKfKd98KT"
      },
      "source": [
        "The ```grad()``` function takes as input a function ```f``` and returns the function ``` f' ``` . This ```f'``` should be ```jit()```-ted again."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JpKDUNP38hGD",
        "outputId": "f51a44c0-51a2-4e21-bbc6-27328acec43b"
      },
      "source": [
        "def FiniteDiffGrad(x):\n",
        "    \"\"\" Compute the finite difference derivative approx for the ReLU\"\"\"\n",
        "    return np.array((ReLU(x + 1e-3) - ReLU(x - 1e-3)) / (2 * 1e-3))\n",
        "\n",
        "# Compare the Jax gradient with a finite difference approximation\n",
        "print(\"Jax Grad: \", jit(grad(jit(ReLU)))(2.))\n",
        "print(\"FD Gradient:\", FiniteDiffGrad(2.))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Jax Grad:  1.0\n",
            "FD Gradient: 0.99998707\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VlG3QboG-dRO"
      },
      "source": [
        "**vmap** -  makes batching as easy as never before. While in PyTorch one always has to be careful over which dimension you want to perform computations, vmap lets you simply write your computations for a single sample case and afterwards wrap it to make it batch compatible. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bfS5MpeR96Jj"
      },
      "source": [
        "batch_dim = 32\n",
        "feature_dim = 100\n",
        "hidden_dim = 512\n",
        "\n",
        "# Generate a batch of vectors to process\n",
        "X = random.normal(key, (batch_dim, feature_dim))\n",
        "\n",
        "# Generate Gaussian weights and biases\n",
        "params = [random.normal(key, (hidden_dim, feature_dim)),\n",
        "          random.normal(key, (hidden_dim, ))]\n",
        "\n",
        "def relu_layer(params, x):\n",
        "    \"\"\" Simple ReLu layer for single sample \"\"\"\n",
        "    return ReLU(np.dot(params[0], x) + params[1])\n",
        "\n",
        "def batch_version_relu_layer(params, x):\n",
        "    \"\"\" Error prone batch version \"\"\"\n",
        "    return ReLU(np.dot(X, params[0].T) + params[1])\n",
        "\n",
        "def vmap_relu_layer(params, x):\n",
        "    \"\"\" vmap version of the ReLU layer \"\"\"\n",
        "    return jit(vmap(relu_layer, in_axes=(None, 0), out_axes=0))\n",
        "\n",
        "out = np.stack([relu_layer(params, X[i, :]) for i in range(X.shape[0])])\n",
        "out = batch_version_relu_layer(params, X)\n",
        "out = vmap_relu_layer(params, X)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bBgq7oJIKGC5"
      },
      "source": [
        "```vmap``` wraps the ```relu_layer``` function and takes as an input the axis over which to batch the inputs. In our case the first input to ```relu_layer``` are the parameters which are the same for the entire batch [```(None)```]. The second input is the feature vector, ```x```. We have stacked the vectors into a matrix such that our input has dimensions ```(batch_dim, feature_dim)```. We therefore need to provide ```vmap``` with batch dimension ```(0)``` in order to properly parallelize the computations. ```out_axes``` then specifies how to stack the individual samples' outputs. In order to keep things consistent, we choose the first dimension to remain the batch dimension."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAlWJJKWMI8q"
      },
      "source": [
        "## MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rTUvezvFKnOl"
      },
      "source": [
        "from jax.scipy.special import logsumexp\n",
        "from jax.experimental import optimizers\n",
        "\n",
        "import torch\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zJrb8CtwMTIg"
      },
      "source": [
        "batch_size = 100\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('../data', train=True, download=True,\n",
        "                   transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=batch_size, shuffle=True)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=batch_size, shuffle=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xZeZBPooOZLH",
        "outputId": "5ac0521b-5c56-4c62-d155-07198de28c3d"
      },
      "source": [
        "print(key)\n",
        "split = random.split(key, 5)                                                    #Can be split into any number of parts. New keys, along new axis\n",
        "print(split)                                                \n",
        "print(random.split(split[0]))                                                   #Can only split \"keys\", i.e. , nd-array of size (2,)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[0 1]\n",
            "[[3243370355 1344208528]\n",
            " [ 532076793 2354449600]\n",
            " [1813813011 1313272271]\n",
            " [3522235465 4107438537]\n",
            " [1531693580 2391939978]]\n",
            "[[1467608531 2825924092]\n",
            " [ 757006082 1868645737]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cqSnsYLyR72N"
      },
      "source": [
        "Since ```JAX``` offers only a functional programming interface, we can't write classes corresponding to modules, in ```JAX``` . We must write a function for initialization, and forward pass instead. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qy3Pytv4NW5g"
      },
      "source": [
        "def initialize_mlp(sizes, key):\n",
        "    \"\"\" Initialize the weights of all layers of a linear layer network \"\"\"\n",
        "\n",
        "    keys = random.split(key, len(sizes))\n",
        "    \n",
        "    # Initialize a single layer with Gaussian weights -  helper function\n",
        "    def initialize_layer(m, n, key, scale=1e-2):\n",
        "        w_key, b_key = random.split(key)\n",
        "        return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n",
        "    \n",
        "    return [initialize_layer(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n",
        "\n",
        "\n",
        "layer_sizes = [784, 512, 512, 10]\n",
        "\n",
        "# Return a list of tuples of layer weights\n",
        "params = initialize_mlp(layer_sizes, key)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MtVhCfKlSnLQ"
      },
      "source": [
        "The forward passs functions should take as input all the parameters(```params```) of the model, and the input(```in_array```) to it. Usually, we make a dictionary of all the parameters, so that the function can access them easily."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yxHYcJ8fQSoJ"
      },
      "source": [
        "def forward_pass(params, in_array):\n",
        "    \"\"\" \n",
        "    Compute the forward pass for each example individually.\n",
        "    Inputs :  params: List of tuples. Tuples must be as required by relu_layer.\n",
        "              in_array: Input array as needed by relu_layer.\n",
        "    \"\"\"\n",
        "    activations = in_array\n",
        "\n",
        "    # Loop over the ReLU hidden layers\n",
        "    for w, b in params[:-1]:\n",
        "        activations = relu_layer([w, b], activations)\n",
        "\n",
        "    # Perform final trafo to logits\n",
        "    final_w, final_b = params[-1]\n",
        "    logits = np.dot(final_w, activations) + final_b                             #Feel free to use any jit-numpy operations in your functions, anywhere.\n",
        "\n",
        "    return logits - logsumexp(logits)                                           #Just simple softmax, it is. \n",
        "\n",
        "# Make a batched version of the `predict` function\n",
        "batch_forward = vmap(forward_pass, in_axes=(None, 0), out_axes=0)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZoCF9Cu2R5Fy"
      },
      "source": [
        "def one_hot(x, k, dtype=np.float32):\n",
        "    \"\"\"Create a one-hot encoding of x of size k \"\"\"\n",
        "    return np.array(x[:, None] == np.arange(k), dtype)\n",
        "\n",
        "def loss(params, in_arrays, targets):\n",
        "    \"\"\" \n",
        "    Compute the multi-class cross-entropy loss.\n",
        "\n",
        "    Inputs : params: list of model parameters as accepted by forward_pass\n",
        "             in_arrays: input_array as accepted by forward_pass\n",
        "             targets: jit-numpy array containing one hot targets\n",
        "    \"\"\"\n",
        "    preds = batch_forward(params, in_arrays)\n",
        "    return -np.sum(preds * targets)                                             #Cross Entropy Loss. Divide by 784 to average.\n",
        "\n",
        "def accuracy(params, data_loader):\n",
        "    \"\"\" Compute the accuracy for a provided dataloader \"\"\"\n",
        "    acc_total = 0\n",
        "    num_classes = 10\n",
        "\n",
        "    for batch_idx, (data, target) in enumerate(data_loader):\n",
        "        images = np.array(data).reshape(data.size(0), 28*28)                    #Need to make PyTorch tensors, into jit-numpy arrays\n",
        "        targets = one_hot(np.array(target), num_classes)\n",
        "\n",
        "        target_class = np.argmax(targets, axis=1)\n",
        "        predicted_class = np.argmax(batch_forward(params, images), axis=1)\n",
        "        acc_total += np.sum(predicted_class == target_class)\n",
        "    return acc_total/len(data_loader.dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OzuAhdz6UQlN",
        "outputId": "9ec82b1d-0077-4501-f0ba-603ca56ce437"
      },
      "source": [
        "x = np.arange(3)\n",
        "print(x.shape)\n",
        "print(x[None, :].shape)\n",
        "print(x[:,None].shape)\n",
        "print(x+x[None,:])\n",
        "print(x[None,:]+x[:,None])\n",
        "print(x+x[:,None])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "(3,)\n",
            "(1, 3)\n",
            "(3, 1)\n",
            "[[0 2 4]]\n",
            "[[0 1 2]\n",
            " [1 2 3]\n",
            " [2 3 4]]\n",
            "[[0 1 2]\n",
            " [1 2 3]\n",
            " [2 3 4]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dzzVCvqpcWqw"
      },
      "source": [
        "```value_and_grad(fn)``` returns a function that takes same arguments(```x```) as ```fn``` and returns both the return value(```fn(x)```) of ```fn``` and its gradient(```fn'(x)```), as a tuple. \n",
        "\n",
        "The optimizer below stores its data(parameters and hyperparameters) in ```opt_state``` and its functionality is defined in ```opt_update()```, ```opt_init()``` and ```get_params()``` . Notice how there is no class. It would be better to put all 4 things in a dicionary, hence."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39Pivo8eUk9X"
      },
      "source": [
        "@jit\n",
        "def update(params, x, y, opt_state):\n",
        "    \"\"\" \n",
        "    Compute the gradient for a batch and update the parameters\n",
        "\n",
        "    Inputs :  params:     list of model parameters as accepted by loss function (in turn by forward_pass)\n",
        "              x:          input as accepted by loss_function(in turn by forward_pass)\n",
        "              y:          jit-numpy array containing one hot targets(as required by loss function)\n",
        "              opt_state:  as required by opt_update\n",
        "    Returns : \n",
        "              updated parameters, current optimizer state, computed value\n",
        "    \"\"\"\n",
        "    value, grads = value_and_grad(loss)(params, x, y)\n",
        "    opt_state = opt_update(0, grads, opt_state)                                 #opt_update is a function, not a variable, hence is available in this scope, although not defined here.\n",
        "    return get_params(opt_state), opt_state, value                              #The first argument to the opt_update function is the optimizer step number.\n",
        "\n",
        "# Defining an optimizer in Jax\n",
        "step_size = 1e-3\n",
        "opt_init, opt_update, get_params = optimizers.adam(step_size)\n",
        "opt_state = opt_init(params)                                                    #All the updatable parameters. First opt_state needs to be obtained this way, always.\n",
        "\n",
        "num_epochs = 10\n",
        "num_classes = 10"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tnZioQ7tlGzx"
      },
      "source": [
        "Notice how in all the above code, each function tries to make sure that its input fits well with the functions that it is calling. And this leads to a hierarchical structure, in stark comparison to the step-wise structure of PyTorch code. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nY3bya5GeQBF",
        "outputId": "0e85d070-9a6b-4784-c3e4-8fe676b4197e"
      },
      "source": [
        "def run_mnist_training_loop(num_epochs, opt_state, net_type=\"MLP\"):\n",
        "    \"\"\" Implements a learning loop over epochs. \"\"\"\n",
        "\n",
        "    # Initialize placeholder for logging\n",
        "    log_acc_train, log_acc_test, train_loss = [], [], []\n",
        "\n",
        "    # Get the initial set of parameters\n",
        "    params = get_params(opt_state)                                              #Assumes all parameters are updatable. Otherwise send as argument in this function. \n",
        "\n",
        "    # Get initial accuracy after random init\n",
        "    train_acc = accuracy(params, train_loader)\n",
        "    test_acc = accuracy(params, test_loader)\n",
        "    log_acc_train.append(train_acc)\n",
        "    log_acc_test.append(test_acc)\n",
        "\n",
        "    # Loop over the training epochs\n",
        "    for epoch in range(num_epochs):\n",
        "        start_time = time.time()\n",
        "        for batch_idx, (data, target) in enumerate(train_loader):\n",
        "            if net_type == \"MLP\":\n",
        "                # Flatten the image into 784-sized vectors for the MLP\n",
        "                x = np.array(data).reshape(data.size(0), 28*28)\n",
        "            elif net_type == \"CNN\":\n",
        "                # No flattening of the input required for the CNN\n",
        "                x = np.array(data)\n",
        "            y = one_hot(np.array(target), num_classes)\n",
        "            params, opt_state, loss = update(params, x, y, opt_state)\n",
        "            train_loss.append(loss)\n",
        "\n",
        "        epoch_time = time.time() - start_time\n",
        "        train_acc = accuracy(params, train_loader)\n",
        "        test_acc = accuracy(params, test_loader)\n",
        "        log_acc_train.append(train_acc)\n",
        "        log_acc_test.append(test_acc)\n",
        "        print(\"Epoch {} | T: {:0.2f} | Train A: {:0.3f} | Test A: {:0.3f}\".format(epoch+1, epoch_time,\n",
        "                                                                    train_acc, test_acc))\n",
        "\n",
        "    return train_loss, log_acc_train, log_acc_test\n",
        "\n",
        "\n",
        "train_loss, train_log, test_log = run_mnist_training_loop(num_epochs,\n",
        "                                                          opt_state,\n",
        "                                                          net_type=\"MLP\")\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1 | T: 16.56 | Train A: 0.973 | Test A: 0.968\n",
            "Epoch 2 | T: 15.61 | Train A: 0.984 | Test A: 0.974\n",
            "Epoch 3 | T: 15.55 | Train A: 0.990 | Test A: 0.979\n",
            "Epoch 4 | T: 15.63 | Train A: 0.993 | Test A: 0.981\n",
            "Epoch 5 | T: 15.41 | Train A: 0.992 | Test A: 0.978\n",
            "Epoch 6 | T: 15.13 | Train A: 0.997 | Test A: 0.982\n",
            "Epoch 7 | T: 15.27 | Train A: 0.996 | Test A: 0.980\n",
            "Epoch 8 | T: 15.93 | Train A: 0.996 | Test A: 0.980\n",
            "Epoch 9 | T: 15.80 | Train A: 0.995 | Test A: 0.981\n",
            "Epoch 10 | T: 15.36 | Train A: 0.997 | Test A: 0.982\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYmvwp9RloXc"
      },
      "source": [
        "# CNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S8h_tPpakkU7"
      },
      "source": [
        "from jax.experimental import stax\n",
        "from jax.experimental.stax import (BatchNorm, Conv, Dense, Flatten,\n",
        "                                   Relu, LogSoftmax)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QOMdDw6hnlsy"
      },
      "source": [
        "The ```init_fun()``` below takes the ```key``` and the shape of input as its arguments. It returns the output shape and the randomly assigned parameters. \n",
        "\n",
        "The ```conv_net()``` function takes ```params``` and input of the shape specified in second argument of ```init_fun()``` and returns the result of the convolution operations specified in ```stax.serial()```. Note that if it is a function that returns ```f(x)``` , you can quickly make another one to get ```f'(x)``` ."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "STD7OZT3l1JT"
      },
      "source": [
        "init_fun, conv_net = stax.serial(Conv(32, (5, 5), (2, 2), padding=\"SAME\"),      #First argument is number of out channels, second is filter shape, third stride.\n",
        "                                 BatchNorm(), Relu,\n",
        "                                 Conv(32, (5, 5), (2, 2), padding=\"SAME\"),\n",
        "                                 BatchNorm(), Relu,\n",
        "                                 Conv(10, (3, 3), (2, 2), padding=\"SAME\"),\n",
        "                                 BatchNorm(), Relu,\n",
        "                                 Conv(10, (3, 3), (2, 2), padding=\"SAME\"), Relu,\n",
        "                                 Flatten,\n",
        "                                 Dense(num_classes),                            #Only final size needs to be specified !! \n",
        "                                 LogSoftmax)\n",
        "\n",
        "output_shape, params = init_fun(key, (batch_size, 1, 28, 28))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U_zUMJ4YqFO5"
      },
      "source": [
        "Various types of initializations can also be specified for each layer. See [here](https://jax.readthedocs.io/en/latest/_modules/jax/experimental/stax.html#serial) for default initializations of each layer. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7kH0B_kXpJ2o"
      },
      "source": [
        "def accuracy(params, data_loader):\n",
        "    \"\"\" Compute the accuracy for the CNN case (no flattening of input)\"\"\"\n",
        "    acc_total = 0\n",
        "    for batch_idx, (data, target) in enumerate(data_loader):\n",
        "        images = np.array(data)\n",
        "        targets = one_hot(np.array(target), num_classes)\n",
        "\n",
        "        target_class = np.argmax(targets, axis=1)\n",
        "        predicted_class = np.argmax(conv_net(params, images), axis=1)\n",
        "        acc_total += np.sum(predicted_class == target_class)\n",
        "    return acc_total/len(data_loader.dataset)\n",
        "\n",
        "def loss(params, images, targets):\n",
        "    preds = conv_net(params, images)\n",
        "    return -np.sum(preds * targets)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DSD2aU40w2Ut"
      },
      "source": [
        "step_size = 1e-3\n",
        "opt_init, opt_update, get_params = optimizers.adam(step_size)\n",
        "opt_state = opt_init(params)\n",
        "num_epochs = 10\n",
        "\n",
        "train_loss, train_log, test_log = run_mnist_training_loop(num_epochs,\n",
        "                                                          opt_state,\n",
        "                                                          net_type=\"CNN\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}