How_to_build_a_JAX_Framework.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "How to build a JAX Framework.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "LQJc4fTutWSY",
        "lbJ7suQRwsbt",
        "Gmleq5II1h1V",
        "nfM1e9G6w7Ze",
        "5z70xwoao83d",
        "9dEWylPRxXiw",
        "8ynP8ZDaWrHO",
        "DQueUYoXJMNj",
        "0bjRXbF4JGaG"
      ],
      "authorship_tag": "ABX9TyNofbWpffaP9EWeQ477njKz",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DanielhCarranza/Curso-Deep-Learning/blob/master/How_to_build_a_JAX_Framework.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdUNYXX3Cip7"
      },
      "source": [
        "# How to build a Jax Framework from Scratch\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HSjs8GezBsG0"
      },
      "source": [
        "import gzip\n",
        "import math\n",
        "import pickle\n",
        "import numpy as onp\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import jax.numpy as np\n",
        "\n",
        "from jax import random\n",
        "from jax import jacfwd, jacrev\n",
        "from jax import grad, jit, vmap\n",
        "from jax.scipy.special import logsumexp\n",
        "from jax.nn import initializers as init\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch import tensor\n",
        "from torch.nn import functional as F\n",
        "\n",
        "from collections import OrderedDict, namedtuple\n",
        "\n",
        "from fastai import datasets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LQJc4fTutWSY"
      },
      "source": [
        "## Dataset: MNIST\n",
        "\n",
        "Comparing Jax with Pytorch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_cz2kCtNtVZy"
      },
      "source": [
        "MNIST_URL='http://deeplearning.net/data/mnist/mnist.pkl'\n",
        "def get_data(URL:str=MNIST_URL, conv_type=tensor):\n",
        "    path = datasets.download_data(MNIST_URL, ext='.gz')\n",
        "    with gzip.open(path, 'rb') as f:\n",
        "        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')\n",
        "    return map(conv_type, (x_train,y_train,x_valid,y_valid))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "89azjfrTtaC2",
        "outputId": "14396e8a-5646-4995-ca60-66709bbd9dd9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "x_train, y_train, x_valid, y_valid = get_data()\n",
        "print(f'X shape: {x_train.shape}, Y shape: {y_valid.shape}')\n",
        "\n",
        "print(f'X mean: {x_train.mean()}, X std: {x_train.std()}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading http://deeplearning.net/data/mnist/mnist.pkl\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              ""
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "X shape: torch.Size([50000, 784]), Y shape: torch.Size([10000])\n",
            "X mean: 0.1304190456867218, X std: 0.30728983879089355\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cRFDCUT3ufpF"
      },
      "source": [
        "# Parameters\n",
        "weights = torch.randn(784,10)\n",
        "bias = torch.zeros(10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jz_EkdQ1wgzO"
      },
      "source": [
        "m1 = x_valid[:5]\n",
        "m2 = weights"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xnWRmw12vQcN",
        "outputId": "e2375a25-dae7-422d-d48d-76be7d0f094e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "# Pytorch\n",
        "%timeit -n 10 t2 = m1 @ m2"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 479.45 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 3: 6.81 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbJ7suQRwsbt"
      },
      "source": [
        "### JAX "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aLk5LlTJw7HO",
        "outputId": "5c797605-0fec-470d-8530-6169693877f3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "x_jtrain, y_jtrain, x_jvalid, y_jvalid = get_data(conv_type=np.array)\n",
        "type(x_jtrain)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "jax.interpreters.xla.DeviceArray"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3vbUEKpIyx1_"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "weights = random.normal(key, (784,10), dtype=np.float32)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YQTUoFiky9vL",
        "outputId": "be306b2a-ff79-41be-e209-16406a2b0d2b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "m1 = x_jvalid[:5]\n",
        "m2 = weights\n",
        "m1.shape, m2.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((5, 784), (784, 10))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4MOOaF8Xw7LC",
        "outputId": "b0f14ee0-92c8-48c5-a369-8678d369b02a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "%timeit -n 10 t2 = m1 @ m2"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 96.64 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 3: 275 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H3H0YTMuw7On",
        "outputId": "3b0fe458-2eba-42d4-ffb7-8b0960b664e6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "%timeit -n 10 t2 = np.matmul(m1, m2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "10 loops, best of 3: 290 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gmleq5II1h1V"
      },
      "source": [
        "### FUNCTION computation comparison\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tXY3fvnMw7Sa",
        "outputId": "d7c11cc9-9fb2-4621-aae4-f84c6240c2ef",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "@jit\n",
        "def selu(x, alpha=1.67, lmbda=1.05):\n",
        "  return lmbda * np.where(x > 0, x, alpha * np.exp(x) - alpha)\n",
        "\n",
        "x = random.normal(key, (1000000,))\n",
        "%timeit selu(x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 848.18 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10000 loops, best of 3: 172 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S1Wo1NtQw7WJ",
        "outputId": "719d6b35-1e9f-46fe-bce7-46cd7a1ffa08",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "def selu(x, alpha=1.67, lmbda=1.05):\n",
        "  return lmbda * torch.where(x > 0, x, alpha * torch.exp(x) - alpha)\n",
        "\n",
        "x = torch.randn(1000000,)\n",
        "%timeit selu(x)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 5.97 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "100 loops, best of 3: 8.44 ms per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nfM1e9G6w7Ze"
      },
      "source": [
        "## NN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nwvWFOReo8ms"
      },
      "source": [
        "@jit\n",
        "def jnormalize(x, m, s):  return (x-m)/s\n",
        "\n",
        "def normalize(x, m, s): return (x-m)/s\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I3XwLkRp5L9c",
        "outputId": "741d1bcc-2746-4c09-b78a-71b84ec695bd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "# PYTORCH\n",
        "train_mean, train_std = x_train.mean(), x_train.std()\n",
        "%timeit -n 10  normalize(x_train, train_mean, train_std)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "10 loops, best of 3: 73.1 ms per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N0VQzppWo8r_",
        "outputId": "6d929b82-fb8a-4927-bbfd-59a0489f2957",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "# JAX\n",
        "train_mean, train_std = x_jtrain.mean(), x_jtrain.std()\n",
        "%timeit -n 10  jnormalize(x_jtrain,train_mean, train_std)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 91.64 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 3: 115 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XLRb4Q-Ko8zg",
        "outputId": "bc7a44ac-78fb-432e-c3c6-ca090b924d70",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "x_jtrain = jnormalize(x_jtrain, train_mean, train_std)\n",
        "x_jvalid = jnormalize(x_jvalid, train_mean, train_std)\n",
        "print(f'MEAN: {train_mean},STD: {train_std}')\n",
        "print(f'MEAN: {x_jtrain.mean()},STD: {x_jtrain.std()}')\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "MEAN: 0.13044890761375427,STD: 0.3072897791862488\n",
            "MEAN: 2.6451584744791035e-06,STD: 0.9999991655349731\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5z70xwoao83d"
      },
      "source": [
        "### Basic Architecture JAX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zj7n4VNuo87c",
        "outputId": "d9cf82ec-afe2-4ba9-fd88-4b2587c8e785",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "nh = 50 # hidden layers\n",
        "n, m = x_jtrain.shape\n",
        "c = int(y_jtrain.max()+1)\n",
        "n,m,c"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(50000, 784, 10)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FOVJasoMo8_G"
      },
      "source": [
        "w1 = random.normal(key, (m, nh), dtype=np.float32)/np.sqrt(m)\n",
        "b1 = np.zeros(nh)\n",
        "w2 = random.normal(key, (nh,1), dtype=np.float32)/np.sqrt(nh)\n",
        "b2 = np.zeros(1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IvgM0kQbo9DP"
      },
      "source": [
        "def linear(x, w, b): return x @ w + b"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "66I_B1tBv3fB",
        "outputId": "674cfd70-d7e4-46bc-f9b9-1b688a9c9971",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "%timeit -n 10 linear(x_jtrain, w1, b1)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 63.46 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 3: 490 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9ENqVUgjo9HS",
        "outputId": "2eb3cf02-ec35-4fec-b397-bed83c13a233",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "t = linear(x_jtrain, w1, b1)\n",
        "t.mean(), t.std()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(0.16632, dtype=float32), DeviceArray(0.987166, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TtK0Fcjwo9K1"
      },
      "source": [
        "@jit\n",
        "def relu(x): return np.maximum(0, x) -0.5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sJ38i-xgwYLs",
        "outputId": "e6e1b171-9062-4389-8bf4-12cfea2b2c17",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "%timeit -n 10 t = relu(linear(x_jvalid, w1, b1))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "10 loops, best of 3: 758 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XPtlMusixXP5",
        "outputId": "a3951e77-bda3-436e-cc88-203ce4e0d23e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "t.mean(), t.std()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(0.166319, dtype=float32), DeviceArray(0.987166, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v_OT9Fc29NDe"
      },
      "source": [
        "### Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aXvAVQ_3ERrk"
      },
      "source": [
        "from jax.nn import initializers as init"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NCX0SBscERwS",
        "outputId": "2de9356c-89da-46db-e337-43a5a2a16723",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "weights= init.kaiming_normal()(key, (m,nh))\n",
        "bias = np.zeros(nh)\n",
        "weights.mean(), weights.std()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(-0.0003, dtype=float32), DeviceArray(0.050342, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NfBIlQCyER0Q",
        "outputId": "9bc1091a-b2ef-454d-b29e-45bd52fcb92f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "t = relu(linear(x_jtrain, weights, bias))\n",
        "t.mean(), t.std()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(0.169347, dtype=float32), DeviceArray(0.894278, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WZOlTIRFTOlC"
      },
      "source": [
        "@jit\n",
        "def model(xb):\n",
        "    l1 = linear(xb, w1, b1)\n",
        "    l2 = relu(l1)\n",
        "    l3 = linear(l2, w2, b2)\n",
        "    return l3"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C-PTk2IMTPV0",
        "outputId": "d1ee7365-dc39-43a5-d20c-0952f53a9905",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "%timeit -n 10 _=model(x_jvalid)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The slowest run took 14.83 times longer than the fastest. This could mean that an intermediate result is being cached.\n",
            "10 loops, best of 3: 250 µs per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9dEWylPRxXiw"
      },
      "source": [
        "### Loss Function: MSE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F3MU8UDsxXmc",
        "outputId": "f286c99e-86d0-4780-d887-752a2ff8bb34",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "model(x_jvalid).shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000, 1)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0liCT-5ExXqI"
      },
      "source": [
        "@jit\n",
        "def mse(output, target): return ((output.squeeze(axis=-1) - target)**2).mean()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hAUcZprCxXt3",
        "outputId": "c003b6b0-8872-45f6-8e38-e4208a63ed08",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "preds = model(x_jtrain)\n",
        "preds.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(50000, 1)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 73
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KEgoENpexX06",
        "outputId": "e24a51aa-af39-4795-e4d7-50fe7fc66bc4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "%time mse(preds, y_jtrain)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 1.14 ms, sys: 0 ns, total: 1.14 ms\n",
            "Wall time: 741 µs\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(28.006317, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7-cIiLFmwb0_",
        "outputId": "9e94f4c9-bbc2-40df-bb5d-e0996cff31c0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        }
      },
      "source": [
        "grad(mse)(preds, y_jtrain)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[-1.855972e-04],\n",
              "             [ 2.692820e-06],\n",
              "             [-1.226115e-04],\n",
              "             [-1.319953e-05],\n",
              "             ...,\n",
              "             [ 1.071363e-05],\n",
              "             [-3.087608e-04],\n",
              "             [-1.495328e-04],\n",
              "             [-3.050688e-04]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 78
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ynP8ZDaWrHO"
      },
      "source": [
        "### Gradients and Backward pass\n",
        "\n",
        "\\begin{align} \n",
        "& \\bullet \\text{Initialize } W^{[1]} .. W^{[L]}, b^{[1]} … b^{[L]} \\\\ \n",
        "& \\bullet \\text{Set } A^{[0]} = X \\text{ ( Input ) }, L = \\text{Total Layers}  \\\\ \n",
        "& \\bullet \\text{Loop } \\text{epoch} = 1 \\text{ to } \\text{ max iteration } \\\\ \n",
        "& \\rule{1cm}{0pt} \\bullet \\text{Forward Propagation} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet \\text{Loop } l=1 \\text{ to } L-1 \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet Z^{[l]} = W^{[l]}A^{[l-1]}+b^{[l]} \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet A^{[l]} = g \\left (  b^{[l]} \\right )  \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet \\text{Save }  A^{[l]},W^{[l]} \\text{ in memory for later use } \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet Z^{[L]} = W^{[L]}A^{[L-1]}+b^{[L]} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet A^{[L]} = \\sigma \\left (  Z^{[L]} \\right )  \\\\ \n",
        "& \\rule{1cm}{0pt} \\bullet \\text{Cost } J= – \\frac{1}{n} \\bigg(  Ylog \\left ( A^{[2]} \\right ) – \\left ( 1-Y  \\right )log \\left  ( 1 – A^{[2]} \\right ) \\bigg)\\\\ \n",
        "& \\rule{1cm}{0pt} \\bullet \\text{Backward Propagation} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet dA^{[L]} = -\\frac{Y}{A^{[L]}}  + \\frac{1-Y}{1- A^{[L]}} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet dZ^{[L]} = dA^{[L]} \\sigma’\\left ( dA^{[L]} \\right )  \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet dW^{[L]} = dZ^{[L]} dA^{[L-1]} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet db^{[L]} = dZ^{[L]} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet dA^{[L-1]} = dZ^{[L]} W^{[L]} \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet \\text{Loop } l=L-1 \\text{ to } 1 \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet dZ^{[l]} = dA^{[l]} g’\\left ( dA^{[l]} \\right )  \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet dW^{[l]} = dZ^{[l]} dA^{[l-1]} \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet db^{[l]} = dZ^{[l]} \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet dA^{[l-1]} = dZ^{[l]} W^{[l]} \\\\ \n",
        "& \\rule{1cm}{0pt} \\bullet \\text{Update W and b}  \\\\ \n",
        "& \\rule{2cm}{0pt} \\bullet \\text{Loop } l=1 \\text{ to } L \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet W^{[l]} =W^{[l]} -\\alpha . dW^{[l]} \\\\ \n",
        "& \\rule{3cm}{0pt} \\bullet b^{[l]} =b^{[l]} -\\alpha . db^{[l]} \n",
        "\\end{align}"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZyplDv7d4wxU"
      },
      "source": [
        "def mse_grad(inp, target):\n",
        "  inp.g = 2 * (inp.squeeze() - target)[...,None]/inp.shape[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hl7EUNI6Lbxo"
      },
      "source": [
        "def relu_grad(inp, out):\n",
        "  inp.g = float(inp>0)*out.g "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r1zx_tYa1-fd"
      },
      "source": [
        "$$\n",
        "d A^{[l-1]}=\\frac{\\partial \\mathscr{L}}{\\partial A^{[l-1]}}=W^{[l] T} d Z^{[l]}\n",
        "$$\n",
        "\n",
        "$$\n",
        "d Z^{[l]}=d A^{[l] *} g^{\\prime}\\left(Z^{[l]}\\right)\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qj3FpqL9Lb1i"
      },
      "source": [
        "def linear_grad(inp, out, w, b):\n",
        "  #grad of matmul with respect to input\n",
        "  inp.g = out.g @ w.T()\n",
        "  w.g = (inp[...,None] * out.g[:,None,:]).sum(axis=0)\n",
        "  b.g = out.g.sum(axis=0, keepdims=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pcuRUfawLb48"
      },
      "source": [
        "def forward_and_backward(inp, target):\n",
        "  # Forward\n",
        "  l1 = linear(inp,w1,b1)\n",
        "  l2 = relu(l1)\n",
        "  out = linear(l2, w2,b2)\n",
        "  mse(out, target)\n",
        "\n",
        "  # Backward\n",
        "  # mse_grad(out, target)\n",
        "  linear_grad(l2, out, w2, b2)\n",
        "  relu_grad(l1, l2)\n",
        "  linear_grad(inp, l1, w1,b1)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z2LfGlwqLb8g"
      },
      "source": [
        "forward_and_backward(x_jtrain,y_jtrain)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9Zy-z4HvLcHO",
        "outputId": "06735359-0994-4e34-8cf1-9de1867cfd18",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "a = x_jtrain[:5]\n",
        "a.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, 784)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NZlRjbo2GoMc"
      },
      "source": [
        "def forward(inp, target):\n",
        "  # Forward\n",
        "  l1 = linear(inp,w1,b1)\n",
        "  l2 = relu(l1)\n",
        "  out = linear(l2, w2,b2)\n",
        "  return mse(out, target)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "skG-6wHyGoRG"
      },
      "source": [
        "loss= forward(x_jtrain, y_jtrain)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Jhou5WljPG5"
      },
      "source": [
        "class Tensorjax(np.DeviceArray):\n",
        "  def __init__(self, data):\n",
        "    self.data = np.DeviceArray(data)\n",
        "  def __repr__(self): print(f'{data.shape}, {data} ')\n",
        "  def gradient(self, x):\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ddKvMpsQj1Aw"
      },
      "source": [
        "# Tensorjax.gradient = 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lk-bnTIhGoUu"
      },
      "source": [
        "## JAX Framework"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2UfMJmI1gXDI"
      },
      "source": [
        "class Module():\n",
        "  def __init__(self): self._modules={}\n",
        "  def __call__(self, *args): \n",
        "    self.args = args\n",
        "    self.out = self.forward(*args)\n",
        "    return self.out\n",
        "  def forward(self): raise Exception(\"Not implemented method\")\n",
        "  \n",
        "  def __setattr__(self, k, v):\n",
        "    if not k.startswith(\"_\"): self._modules[k]=v\n",
        "    super().__setattr__(k,v)\n",
        " \n",
        "  def __repr__(self): return f'{self._modules}'\n",
        " \n",
        "  def parameters(self):\n",
        "    for l in self._modules.values():\n",
        "      for p in l.parameters: yield p\n",
        "\n",
        "      "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i-2-HHRPYrzf"
      },
      "source": [
        "# class Module():\n",
        "#   def __call__(self, *args): \n",
        "#     self.args = args\n",
        "#     self.out = self.forward(*args)\n",
        "#     return self.out\n",
        "#   def forward(self): raise Exception(\"Not implemented method\")\n",
        "\n",
        "class ReLU(Module):\n",
        "  def forward(self, x):\n",
        "    return np.maximum(0,x) - 0.5\n",
        "\n",
        "class Sigmoid(Module): \n",
        "  def forward(self, x):\n",
        "    return 0.5*(np.tanh(x/2) + 1)\n",
        "\n",
        "class Linear(Module):\n",
        "  def __init__(self, W, b):\n",
        "    super().__init__() \n",
        "    self.W, self.b = W, b\n",
        "  def forward(self, inputs):\n",
        "    return (inputs @ self.W + self.b)\n",
        "\n",
        "\n",
        "class LinearV2(Module):\n",
        "  def forward(self, W, b, inputs):\n",
        "    return (inputs @ W + b)\n",
        "\n",
        "class Loss(Module):\n",
        "  def __init__(self, inputs, targets):\n",
        "    super().__init__()\n",
        "    self.x, self.y = inputs, targets\n",
        "\n",
        "  def forward(self, W, b):\n",
        "    preds = Sigmoid()(LinearV2()(W, b, self.x))\n",
        "    label_probs = preds*self.y + (1-preds)*(1-self.y)\n",
        "    return -np.sum(np.log(label_probs))\n",
        "\n",
        "def cross_entropy(preds, targets):\n",
        "    label_probs = preds*targets + (1-preds)*(1-targets)\n",
        "    return -np.sum(np.log(label_probs)) \n",
        "\n",
        "class LossV2(Module):\n",
        "  def __init__(self, model, inp):\n",
        "    self.model, self.inp = model, inp\n",
        "  def forward(self, params, preds, targets):\n",
        "    preds = self.model(params, self.inp)\n",
        "    return cross_entropy(preds, targets)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sa04rLldbpL8",
        "outputId": "df49948a-4abb-46cd-d8ef-0334498687a7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "# Build a toy dataset.\n",
        "inputs = np.array([[0.52, 1.12,  0.77],\n",
        "                   [0.88, -1.08, 0.15],\n",
        "                   [0.52, 0.06, -1.30],\n",
        "                   [0.74, -2.49, 1.39]])\n",
        "targets = np.array([True, True, False, True])\n",
        "\n",
        "\n",
        "# Initialize random model coefficients\n",
        "key, W_key, b_key = random.split(key, 3)\n",
        "W = random.normal(W_key, (3,))\n",
        "b = random.normal(b_key, ())\n",
        "\n",
        "grad(Loss(inputs, targets), argnums=(0,1))(W, b)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray([-0.169656, -0.877464, -1.490134], dtype=float32),\n",
              " DeviceArray(-0.292272, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eQlABDQjO-L3",
        "outputId": "c95a6e24-a1ca-4476-b4a1-3856db817c49",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "preds = Sigmoid()(Linear(W,b)(inputs))\n",
        "print(preds, preds.shape == targets.shape)\n",
        "\n",
        "preds = Sigmoid()(LinearV2()(W,b, inputs))\n",
        "print(preds, preds.shape == targets.shape)\n",
        "# grad(LossV2(preds))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[0.758798 0.761036 0.038418 0.994809] True\n",
            "[0.758798 0.761036 0.038418 0.994809] True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VYRiTeDaUYnO",
        "outputId": "65a33f1a-e3a5-4361-96d2-7a3a14da8b50",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "pred = Linear(W,b)\n",
        "pred"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'W': DeviceArray([-0.09006 , -0.718414,  2.477131], dtype=float32), 'b': DeviceArray(0.090163, dtype=float32)}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EcPZK91leBvD"
      },
      "source": [
        "class Linear(Module):\n",
        "  def __init__(self, params): self.params = params\n",
        "  def forward(self, inputs):\n",
        "    return (inputs @ self.params['W'] + self.params['b'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fjN_txnIhtaY"
      },
      "source": [
        "def compute_loss(params, x, targets):\n",
        "  layers = [Linear(params), Sigmoid()]\n",
        "  for l in layers: x = l(x) \n",
        "  return cross_entropy(x, targets)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "knIYag2iihxx",
        "outputId": "c2583cd7-7acc-4e86-c5da-01ff6ebd185d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "grad(compute_loss)({'W':W, 'b':b}, inputs, targets)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'W': DeviceArray([-1.72562 ,  2.016514, -2.178545], dtype=float32),\n",
              " 'b': DeviceArray(-2.381826, dtype=float32)}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 146
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PFFkCLLTPw50",
        "outputId": "abe35a98-bf33-4ded-8854-3366140236cc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "nh = 50 # hidden layers\n",
        "n, m = x_jtrain.shape\n",
        "c = int(y_jtrain.max()+1)\n",
        "n,m,c"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(50000, 784, 10)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sz2dawLO3Fi9"
      },
      "source": [
        "def init_params(n_in, n_out, init=init.kaiming_normal(), key=random.PRNGKey(0)):\n",
        "  W_key, b_key = random.split(key, 2)\n",
        "  weights= init(W_key, (n_in,n_out))\n",
        "  bias = np.zeros(n_out)\n",
        "  return weights, bias "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJbklDj12-O7"
      },
      "source": [
        "w1, b1 = init_params(m, nh)\n",
        "w2, b2 = init_params(nh, 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aqfiSiaj5N4u",
        "outputId": "6fcd3865-a850-4fb0-c5d4-156c918c8229",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "t = ReLU()(Linear(w1, b1)(x_jtrain))\n",
        "t.mean(), t.std()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(-0.352507, dtype=float32), DeviceArray(0.236954, dtype=float32))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DQueUYoXJMNj"
      },
      "source": [
        "#### Jax models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MZrw21fel1TO"
      },
      "source": [
        "def log_softmax(logits):\n",
        "  return logits - logsumexp(logits)[..., None]\n",
        "\n",
        "def cross_entropy(preds, targets):\n",
        "  return -np.sum(log_softmax(preds) * targets)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z_MMGiXGSgET"
      },
      "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 model(x,y):\n",
        "  layers = [Linear(w1,b1), ReLU(), Linear(w2,b2)]\n",
        "  for l in layers: x = l(x)\n",
        "  return x "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lR6B-EZFELi5"
      },
      "source": [
        "y = one_hot(y_jtrain, c)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pA6xUT1f2pCr",
        "outputId": "26af791b-e43c-48ab-9903-95d334343204",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "%timeit -n 10 _ = model(x_jtrain, y)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "10 loops, best of 3: 1.99 ms per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qS1GRB5v2sU6",
        "outputId": "01625ba8-e7a7-4564-d251-68b4f0ec10f7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        }
      },
      "source": [
        "def compute_loss( x,y):\n",
        "  preds= model(x,y)\n",
        "  return cross_entropy(preds, y)\n",
        "\n",
        "%time compute_loss(x_jtrain, y)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 5.03 ms, sys: 3.46 ms, total: 8.5 ms\n",
            "Wall time: 10.7 ms\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(543525.4, dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08cYcA6KITT7"
      },
      "source": [
        "### JAX Module"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KPhoVxRxjWj1"
      },
      "source": [
        "class Module():\n",
        "  def __init__(self):\n",
        "    self._modules = OrderedDict()\n",
        "    self._parameters = OrderedDict()\n",
        "\n",
        "  def __call__(self, *args): \n",
        "    self.args = args\n",
        "    self.out = self.forward(*args)\n",
        "    return self.out\n",
        "\n",
        "  def forward(self): raise Exception(\"Not implemented method\")\n",
        "  \n",
        "  def __setattr__(self, k, v):\n",
        "    if not k.startswith(\"_\"): self._modules[k]=v\n",
        "    super().__setattr__(k,v)\n",
        "\n",
        "  def __getattr__(self, name):\n",
        "        if '_parameters' in self.__dict__:\n",
        "            _parameters = self.__dict__['_parameters']\n",
        "            if name in _parameters:\n",
        "                return _parameters[name]\n",
        "\n",
        "        if '_modules' in self.__dict__:\n",
        "            modules = self.__dict__['_modules']\n",
        "            if name in modules:\n",
        "                return modules[name]\n",
        "        raise AttributeError(\"'{}' object has no attribute '{}'\".format(\n",
        "            type(self).__name__, name))\n",
        " \n",
        "  def __repr__(self): return f'{self._modules}'\n",
        "\n",
        "  def register_parameter(self, name, param): \n",
        "    if '_parameters' not in self.__dict__:\n",
        "            raise AttributeError(\n",
        "                \"cannot assign parameter before Module.__init__() call\")\n",
        "    self._parameters[name] = param\n",
        "\n",
        "  def parameters(self):\n",
        "    for l in self._modules.values():\n",
        "      for p in l.parameters: yield p"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RRKV8B7wITP1"
      },
      "source": [
        "class Model(Module):\n",
        "  def __init__(self, n_in, nh, n_out):\n",
        "    super().__init__()\n",
        "    self.l1 = nn.Linear(n_in, nh)\n",
        "    self.l2 = nn.Linear(nh, n_out)\n",
        "  \n",
        "  def forward(self, x):\n",
        "    return self.l2(nn.ReLU((self.l2(x))))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4HWTvYbS7r5Y",
        "outputId": "a1350d1a-4522-4a59-c96d-00c700ab6bf9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "nh = 50 # hidden layers\n",
        "n, m = x_jtrain.shape\n",
        "c = int(y_jtrain.max()+1)\n",
        "n,m,c"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(50000, 784, 10)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u7Yc0wgQis2g",
        "outputId": "57921213-a6d1-4d08-c41d-73f7f8bd8782",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "mo = Model(m, nh, c)\n",
        "mo"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "OrderedDict([('l1', Linear(in_features=784, out_features=50, bias=True)), ('l2', Linear(in_features=50, out_features=10, bias=True))])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "liEZNxOOjClN",
        "outputId": "b006f5cf-e117-4cae-ffbd-f0e9afa66572",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "mo.parameters"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<bound method Module.parameters of OrderedDict([('l1', Linear(in_features=784, out_features=50, bias=True)), ('l2', Linear(in_features=50, out_features=10, bias=True))])>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u5HpV3WEaiMg"
      },
      "source": [
        "class Model(nn.Module):\n",
        "  def __init__(self, n_in, nh, n_out):\n",
        "    super().__init__()\n",
        "    self.l1 = nn.Linear(n_in, nh)\n",
        "    self.l2 = nn.Linear(nh, n_out)\n",
        "  \n",
        "  def forward(self, x):\n",
        "    return self.l2(nn.ReLU((self.l2(x))))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F_unXGmKalVY",
        "outputId": "c561c556-39d7-4adb-aee7-8748a4e6f35a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        }
      },
      "source": [
        "mo = Model(m, nh, c)\n",
        "mo"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Model(\n",
              "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
              "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kfWJEEdKalaq",
        "outputId": "dc3e122b-6994-4581-f96c-0a5aa8fe9f12",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        }
      },
      "source": [
        "mo.parameters"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<bound method Module.parameters of Model(\n",
              "  (l1): Linear(in_features=784, out_features=50, bias=True)\n",
              "  (l2): Linear(in_features=50, out_features=10, bias=True)\n",
              ")>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 54
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZWCfAmazc669"
      },
      "source": [
        "#### JAX PARAMS"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6OLFsfPSc6va"
      },
      "source": [
        "def init_params(n_in, n_out, init=init.kaiming_normal(), key=random.PRNGKey(0)):\n",
        "  W_key, b_key = random.split(key, 2)\n",
        "  weights= init(W_key, (n_in,n_out))\n",
        "  bias = np.zeros(n_out)\n",
        "  return weights, bias"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JKEj0vJ2c-O6"
      },
      "source": [
        "nn.Parameter??"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Z-chYca9ori"
      },
      "source": [
        "from jax.interpreters.xla import  DeviceArray\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fMS1G6eBc-Tw"
      },
      "source": [
        "class Parameter(onp.ndarray, metaclass=_ArrayMeta):\n",
        "  def __new__(cls, data=None, *args):\n",
        "    # data=np.asarray(data).view(cls)\n",
        "    return super(Parameter, cls).__new__(cls,data, *args)\n",
        "\n",
        "  def __repr__(self):\n",
        "    return 'Parameter containing:\\n' + super(Parameter, self).__repr__()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a9AvvY-m_Shn"
      },
      "source": [
        "type(Parameter(w1)), Parameter(w1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l-U12O44AjEy"
      },
      "source": [
        "class Parameter(onp.ndarray):\n",
        "    def __new__(cls, x, *args, **kwargs): \n",
        "        return super().__new__(cls, x, *args, **kwargs) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5BHQrT88QOD1",
        "outputId": "c2fc914d-9d87-4101-b03d-7612a5854565",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "p = Parameter((3,))\n",
        "isinstance(p, Parameter)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 270
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eMhZpUiv10yH"
      },
      "source": [
        "from jax.lax import convert_element_type\n",
        "convert_element_type?"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8a0kn1Y32MPM",
        "outputId": "ca1caaa3-0678-46aa-88b2-fdb5dd307a90",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "a = onp.array([1.0,2.0,3.0])\n",
        "convert_element_type(a, new_dtype=np.float32)\n",
        "# np.array"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([1., 2., 3.])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 300
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y5npdedYruxb"
      },
      "source": [
        "class InfoArray(onp.ndarray):\n",
        "\n",
        "    def __new__(subtype, shape, dtype=float, buffer=None, offset=0,\n",
        "                strides=None, order=None, info=None):\n",
        "        # Create the ndarray instance of our type, given the usual\n",
        "        # ndarray input arguments.  This will call the standard\n",
        "        # ndarray constructor, but return an object of our type.\n",
        "        # It also triggers a call to InfoArray.__array_finalize__\n",
        "        return super(InfoArray, subtype).__new__(subtype, shape, dtype,\n",
        "                                                buffer, offset, strides,\n",
        "                                                order)\n",
        "         "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mjB155jLpFOB"
      },
      "source": [
        "obj = InfoArray(shape=(3,), )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JFbKQdGceTni"
      },
      "source": [
        "w1, b1 = init_params(m, nh)\n",
        "w2, b2 = init_params(nh, 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AhIWEcxKc-bf",
        "outputId": "a0318356-7ec0-402e-ade6-60bd010bcd3d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "type(w1)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "jax.interpreters.xla.DeviceArray"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ZWsjhQHalmi"
      },
      "source": [
        "### Linear"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZiMnRwDHqEA-"
      },
      "source": [
        "class Model(Module):\n",
        "  def __init__(self, params):\n",
        "    self.layers = [Linear(w1,b1), ReLU(), Linear(w2,b2)]\n",
        "    self.loss = cross_entropy\n",
        "\n",
        "  def forward(self, x, y):\n",
        "    for l in self.layers: x = l(x)\n",
        "    if hasattr(l, 'weight'):\n",
        "      \n",
        "    return self.loss(x, y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TXYdkkVzkkv"
      },
      "source": [
        "def linear(inp, weight, bias=None):\n",
        "  if inp.size()==2 and bias is not None:\n",
        "    ret = inp @ weight + bias\n",
        "  else:\n",
        "    output = inp @ weight.T\n",
        "    if bias is not None: output += bias\n",
        "    ret = output\n",
        "  return ret\n",
        "\n",
        "class Linear(Module):\n",
        "  def __init__(self, in_features, out_features, bias=True): \n",
        "    super().__init__()\n",
        "    self.weight, self.bias =  init_params(in_features, out_features)\n",
        "  def forward(self, inputs):\n",
        "    return linear(inputs, self.weight, self.bias )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YSmM6VJpPTM1",
        "outputId": "d07aa1a0-eabe-4fa9-d31d-a2d86d13c6e4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 207
        }
      },
      "source": [
        "l1 = Linear(m,nh)\n",
        "l1"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'weight': DeviceArray([[-0.049711, -0.035382, -0.054953, -0.001373, ..., -0.021679, -0.109848, -0.020343, -0.039765],\n",
              "             [ 0.00608 ,  0.002496,  0.020257, -0.107654, ...,  0.030652,  0.010056,  0.00069 ,  0.016928],\n",
              "             [ 0.08513 ,  0.05751 , -0.005576, -0.059141, ..., -0.057368, -0.066713,  0.042022,  0.050596],\n",
              "             [ 0.035249, -0.006973, -0.065923, -0.094957, ...,  0.066852, -0.022646,  0.082308, -0.028639],\n",
              "             ...,\n",
              "             [ 0.03995 ,  0.041362, -0.084902, -0.015864, ..., -0.00484 ,  0.048652,  0.010824,  0.015427],\n",
              "             [-0.012211, -0.05021 , -0.039358,  0.058195, ..., -0.039469,  0.074852,  0.073005,  0.024522],\n",
              "             [ 0.034812,  0.076179, -0.075141, -0.097438, ..., -0.09212 , -0.078528,  0.021074, -0.090035],\n",
              "             [ 0.027886, -0.005652, -0.025512, -0.017421, ...,  0.006612,  0.004661, -0.097214, -0.013013]],            dtype=float32), 'bias': DeviceArray([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
              "             0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],            dtype=float32)}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h_A5aTu5PTAB",
        "outputId": "48e801f9-f5ad-44bb-f553-df0d9d0048c4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "l2 = nn.Linear(m,nh)\n",
        "l2.parameters"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<bound method Module.parameters of Linear(in_features=784, out_features=50, bias=True)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f6e_pC4Tu4Z1"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}