Jax_GD.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Jax.GD.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rahuldave/LearningJax/blob/main/Jax_GD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AlHg1tBH7eiy"
      },
      "source": [
        "# Machine Learning using Jax\n",
        "\n",
        "This notebook is taken from \n",
        "_the [solving y=mx+b... with jax on a tpu pod slice](http://matpalm.com/blog/ymxb_pod_slice/) blog series_\n",
        "\n",
        "Here we will fit a linear regression model using MSE. We'll do this both on a single device using `vmap` as well as across `num_replicas` devices. We'll make some sharding of the input data here in addition to the output data, so we can support a seamless Gradient Descent loop."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cat /proc/cpuinfo"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pDonUX0td8on",
        "outputId": "8a7944ea-4bae-4cef-e2fe-649ff47a9d80"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "processor\t: 0\n",
            "vendor_id\t: GenuineIntel\n",
            "cpu family\t: 6\n",
            "model\t\t: 79\n",
            "model name\t: Intel(R) Xeon(R) CPU @ 2.20GHz\n",
            "stepping\t: 0\n",
            "microcode\t: 0x1\n",
            "cpu MHz\t\t: 2199.998\n",
            "cache size\t: 56320 KB\n",
            "physical id\t: 0\n",
            "siblings\t: 2\n",
            "core id\t\t: 0\n",
            "cpu cores\t: 1\n",
            "apicid\t\t: 0\n",
            "initial apicid\t: 0\n",
            "fpu\t\t: yes\n",
            "fpu_exception\t: yes\n",
            "cpuid level\t: 13\n",
            "wp\t\t: yes\n",
            "flags\t\t: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities\n",
            "bugs\t\t: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa\n",
            "bogomips\t: 4399.99\n",
            "clflush size\t: 64\n",
            "cache_alignment\t: 64\n",
            "address sizes\t: 46 bits physical, 48 bits virtual\n",
            "power management:\n",
            "\n",
            "processor\t: 1\n",
            "vendor_id\t: GenuineIntel\n",
            "cpu family\t: 6\n",
            "model\t\t: 79\n",
            "model name\t: Intel(R) Xeon(R) CPU @ 2.20GHz\n",
            "stepping\t: 0\n",
            "microcode\t: 0x1\n",
            "cpu MHz\t\t: 2199.998\n",
            "cache size\t: 56320 KB\n",
            "physical id\t: 0\n",
            "siblings\t: 2\n",
            "core id\t\t: 0\n",
            "cpu cores\t: 1\n",
            "apicid\t\t: 1\n",
            "initial apicid\t: 1\n",
            "fpu\t\t: yes\n",
            "fpu_exception\t: yes\n",
            "cpuid level\t: 13\n",
            "wp\t\t: yes\n",
            "flags\t\t: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities\n",
            "bugs\t\t: cpu_meltdown spectre_v1 spectre_v2 spec_store_bypass l1tf mds swapgs taa\n",
            "bogomips\t: 4399.99\n",
            "clflush size\t: 64\n",
            "cache_alignment\t: 64\n",
            "address sizes\t: 46 bits physical, 48 bits virtual\n",
            "power management:\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dvx44mUvnzY0"
      },
      "source": [
        "setup stuff for runtime TPU or CPU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hjkvme7eYYZx"
      },
      "source": [
        "USE_TPU = True\n",
        "\n",
        "if USE_TPU:\n",
        "  import jax\n",
        "  import jax.tools.colab_tpu\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "else:\n",
        "  # x8 cpu devices  \n",
        "  import os\n",
        "  os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=2'"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BvJbOLW4YccB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8d5210f6-e3ad-4584-c6bd-9f767af73c8f"
      },
      "source": [
        "import jax\n",
        "print(repr(jax.devices()))\n",
        "num_replicas = len(jax.devices())"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4YFomemlh01O"
      },
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "from jax import make_jaxpr"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PFPE54T4rGDL"
      },
      "source": [
        "## 1: Generate some data\n",
        "\n",
        "Let us generate some training data first. We'll add in some noise"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3cWNrO6_-50b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        },
        "outputId": "7edba706-cc49-4c4f-94e5-d84eed1b0ac1"
      },
      "source": [
        "def y(x, m, b):\n",
        "  return m * x + b\n",
        "\n",
        "datasize = 128\n",
        "x = np.linspace(-10, 10, num=datasize)\n",
        "np.random.shuffle(x)\n",
        "\n",
        "m_true = -1.5\n",
        "b_true = 2.0\n",
        "\n",
        "y_true = y(x, m_true, b_true)\n",
        "\n",
        "noise = ((np.random.rand(*y_true.shape) * 2) - 1) * 3\n",
        "y_true += noise\n",
        "\n",
        "plt.plot(x, y_true, 'o')\n",
        "plt.show()"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXkAAAD6CAYAAABEUDf/AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3df5BdZXkH8O83mwtuwGFDSTEs0ETHQkVKIhlGG9sBtASxAwEr0OlYWp1G29IRh1JDmVFs/yAtQ3H6S5tWRttxNGlRiIKNYLBMnfpjYxJIhJSg0OYaYS0sWrPCZvfpH/fc5ezZ9z0/7vlx7zn3+5nZyd17z95z9u7Nc977nOd9XpoZRESkmZb0+wBERKQ8CvIiIg2mIC8i0mAK8iIiDaYgLyLSYAryIiINljvIkzyD5EMkv0PyAMn3B/efTPIBkk8E/y7Pf7giIpIF89bJk1wJYKWZfZvkKwHsBrARwG8DeM7MtpDcDGC5mX0w7rlOOeUUW7VqVa7jEREZNrt37/6hma1wPbY075Ob2REAR4LbPyb5GIBxAFcAuDDY7FMAvgogNsivWrUKExMTeQ9JRGSokHza91ihOXmSqwCsBfANAKcGJwAA+AGAUz0/s4nkBMmJycnJIg9HRGToFRbkSZ4I4G4AN5jZj8KPWScn5MwLmdlWM1tnZutWrHB+2hARkR4VEuRJttAJ8J82s88Fdz8T5Ou7eftni9iXiIikV0R1DQF8AsBjZvaXoYd2ALguuH0dgHvz7ktERLLJfeEVwHoA7wLwKMm9wX1/AmALgO0k3wPgaQBXF7AvERHJoIjqmv8AQM/Db8n7/Gncs6eN23cexPenpnHa2Chu2nAWNq4dr2LXIiIDrYiRfF/ds6eNmz/3KKZnZgEA7alp3Py5RwHAG+h1UhCRYVH7IH/7zoPzAb5remYWt+88uCBwdwN7e2oaxMulPmlOCiIidVX73jXfn5pOvL872m8H90VrObsnBRGRpql9kD9tbDTxftdoP8p3shARqbPaB/mbNpyF0dbIgvtGWyO4acNZ89+nCeC+k4WISJ3VPiffzaPHXUg9bWx0PlXjEj4p6KKsiDRJ7YM80An0cYH4pg1nLajAATB/8XU8FMh7qdQRERlkjQjySdKM9ruPuyp1bty+b8HziIjUxVAEeSB5tA/4c/ezZhrRi0gt1f7Ca5HiLr6qzFJE6khBPsRVqROmMksRqRsF+ZCNa8dx21XnYoTuVjwqsxSRulGQj9i4dhx3XH1eYu29iEgdDM2F1yzSVuOIiAw6BXmPNNU4IiKDrvFBXjNYRWSYNTrIawariAy7Rgf5tL3mk+jTgIjUVaODfJpe80nK+jSgE4eIVKHRJZRpes0nifs00KvwIiaGl08c9+xp9/ycIiIuhQR5kneRfJbk/tB9t5Jsk9wbfF1WxL6ySNNrPkkRnwaiyjhxiIi4FDWS/ySASx3332lma4Kv+wvaV2rdGazjY6MgOm2Fb7vq3ExpEd+ofwnZ88i7jBOHiIhLITl5M3uY5KoinqtoeevdXb3ogXydKX2LmJw02sL6LbuUpxeRwpSdk7+e5CNBOmd5yfsqRVw/m15TLK40UmsJ8ZOXjilPLyKFKjPIfwzAawCsAXAEwB2ujUhuIjlBcmJycjLXDu/Z08b6LbuwevN9WL9lV2EBcuPaccyZOR/rJcXiSiOd+IqlmJlduA/l6UUkL5oneGV+ok665otm9vosj4WtW7fOJiYmetp/tNSxa2y0hVsvPyd32mP9ll3OFMv42Ci+tvniXM8NAKs33wfXX4IAvrfl7bmfX0Sai+RuM1vneqy0kTzJlaFvrwSw37dtEVwVKwAwNT1TSNqjiEodl+6nD9+pVu2NRSSPQi68kvwMgAsBnELyMIAPA7iQ5Bp01st+CsB7i9iXT1zapJdZrlFldKb0ffroUntjEcmrsHRNEfKka3zplK4i0x5ZZ6v6to875nFV14hISnHpmsa0NfCVOnYVlfbI2uYgbnvfpw8C83l+tT8QkTwa09agW7GyfFlr0WNFpj2yzlaN2z6p7YLaH4hIXo0J8kAn0O/50CX46DVrcs1yjZN1tmrc/UkXc9X+QETyaky6JiztLNduKqQ9NY0RErNmiblw32zVuFG5b/uki7lqfyAieTUyyKcRzZXPBhegk3Lsrtx/XDooafu4E1LWE4qISFSj0jVZ+OrqgfiUSJamZ91PCtMzs/NtEbKkj3qpzS9r1q+I1NPQjuSTUh5xj4dH391A/oFtexekW1yfFLoBOu31gay1+VruUESiGlMnn1VSXX2adgW+yUxLCMx5Xtai2iC4+H6nERJzZirBFGmovrQ1GHSuVEhX2pJLX8rHF+CBci+a+p571kwlmCJDamiDfDi3DqCnnHkvAbvMi6ZpnlslmCLDZWhz8kD+BUV81S8+ZfeiSZr126USTJHhMbQj+SLEpXyiipqUFVc9E638cS10AqgEU2SYDPVIPq9uwL51xwFMTc84txltjRQW3D/yhQN4/ujL+3FVz0Qrf9LW9KtHjkgzaSSf08a149j74ZdbKQC95ffjdIN1OMB3FVHTrx45Is01tCWUdVJ2G+WyV70SkXINRavhJku6UHrSaAvrt+zqOdWiHjkizaV0TQ3EXShtLSF+8tKxXKmWpJbHIlJfCvI14KviGRtt4cRXLMXM7MKUW9Za+LLWrxWR/lO6pgbietis3nyf82eypFp8zw8gVxpIRPpPQb4mfBO3impHHH1+NTsTaQala2qurFSLVqUSaYZCRvIk7wLwawCeNbPXB/edDGAbgFUAngJwtZk9X8T+mi7LxKSs7YjTUsWNSDMUla75JIC/AfBPofs2A/iKmW0huTn4/oMF7a+2kgJ4L2mSvD14XLQqlUgzFJKuMbOHATwXufsKAJ8Kbn8KwMYi9lVnaWaWDkqaRBU3Is1QZk7+VDM7Etz+AYBTXRuR3ERyguTE5ORkiYfTf2kC+KCkSbIscygig6uS6hozM5LO/glmthXAVqDT1qCK4+mXNAG8ijRJ2py/Kw2kRmYi9VLmSP4ZkisBIPj32RL3VQtpZpaWnSbJ04xMjcxE6qfMIL8DwHXB7esA3FvivmohTQAvO02SJ+ffy8/G9b8XkfIVVUL5GQAXAjiF5GEAHwawBcB2ku8B8DSAq4vYV52lLXdMWy0TlzrxPZYn55/0s9F9XnT2Cty9u60JVSJ9pFbDNeVbEOS2q84FAO9jt+882HNbYV9L4hESs2YggPC7Kfp9ln2JSHpxrYY147Wm4lIncY/lyfn7GqXNBgOFaED3DR/iPjUovSNSLPWuqale0i7fn5rONUM2+rNLghF8Vr4L0OqXI1I8BfmaSiq1dD0WXVzkzmvWZA6e4esFvg6YceI+NcR9AlGQF+mN0jU1FZd2cT1WxOIiUVlr95MqhQZlIphIk2gkX1Np0i7hx46+dGzRQuB5R8k3bThr0QVeHwKLLrZGq3FOGm1hanrxYuXqlyPSOwX5GosrtYw+VsTiIq59AMknE2BxoHbl31sjRGsJMTP3cp5f/XJE8lGQHxJltUtIWmwEcAdqV/59ZtawfFkLy45bqrYJIgVRkB8SrtRKGaPktNU7vk8QU0dnsOdDl6Ten3rpiMRTkB8SVa7jmmbGbhGfLFRyKZJMM16HmCu10lpCnPiKpZg6OlPqyDhuxm7a/flm4GpGrQybuBmvGskPMWdefM7mL5yWOTIuYtlClVyKJFOQH2JpgmGZk5HyLluoJQpFkinIDzFfkIxynQz6ccGzu8/21LS3KZpKLkUW0ozXIeZrOBblq3GvcvGQ8D6BhU3RGGyjJQpFFtNIfohF8+Injbbwk5eOYWY2fjJSP3rMuPbZZYi/2KoySxlmCvJDzjWZqdca9zIveCY9d3tq2lkK6iuznHj6OTz0+KQCvzSegrwsUFWNe1ZJ1w+IlztvhquCfJ86Pv31/57P5au+XppMOXnJrOzFxsO6i4i0p6bnc+9RrhWouukj3ycA3/YiTaMgL4miqzUBKHWx8fB+wxdbwxdZR9i5NT42GrsCVZZPF6qvlyZSukZi+XLat111bumzSl2pFtdFVt/M126uPTqz1rf2rOrrpYlKH8mTfIrkoyT3klTPgpqJq6QpW9oLvHHpo41rxxd96vjNN55ZWbpJpN+qGslfZGY/rGhfUqB+tg5Ie4E3qUWC62Lyup87WWWVMhSUrpFYRVXS9FKrnqU9ctYWCXlbKojURRUXXg3Al0nuJrkp+iDJTSQnSE5MTk5WcDiSRRGVNL3OkHWlWjSjVSSb0lsNkxw3szbJnwXwAIA/NLOHXduq1fBgyjtj1HdhdITEHVefp6AtklNfWw2bWTv491mSnwdwAQBnkJfBlDe14cvfz5rVZhKSWiNIXZWariF5AslXdm8DuATA/jL3KYMnLn9fh0lI/WjIJlKUsnPypwL4D5L7AHwTwH1m9m8l71MGTFK3yyonIUUndqUJ1P0sIxXJq9R0jZl9F8B5Ze5DBl83rXHj9n3zLYLDxpa1Cl9n1qXXNWF9J6H21PT8SUKpHBlUKqGUSnSD3qI1ZUeI//vpsUVLDpbRJbLXFslxzdFu+pd9ADHfnlnNzmTQqHeNVMZVEnnCcUsxM7dwdN/tEll0DrzXiV1x6aaZOVvQfx9QKkcGi0byUqlopc7qzfc5t/N1iUwaHcdVwfQ6sav78zds2xu7XZiancmg0Ehe+qrILpFJVTB5JnZtXDuO8QzHqmZnMigU5KWvXIHX1zc+KXAmVcHknUHrOtbWEqI1svCIuwuYpK3eESmT0jXSV67mYhedvQJ3726n6lkTlibn7lruMG1lj68RWve+7sImWnFKBknpbQ2yUFsD6eplhqmvfYJvke9oSSXQOZm84/zxnip7su5fpCh9bWsg0oteWilk6VoJFL/+az/bMov4KCcvjZE15170+q++awZLSOXmpW80kpdGyfIJIG6SU1Sa0bjrkwRQr0Zs0jwaycvQKrKyB3j5k0R3kfEwTZCSflGQl6FVxvqvG9eOY85TzKDcvPSD0jUy1MpY/7WoJRNFiqAgLxKRd5GUrFU+ImVSkBcpmG/SVPTEodWmpAoK8tJY/QyiSZ8Geu1tL5KVLrxKIw36kn1abUqqoiAvjTToQVSzY6UqCvLSSIMeRH2VNqrAkaIpyEsjDXoQzdPbXiSL0oM8yUtJHiR5iOTmsvcnAgx+EM3b296l2zZ59eb71Mte5pVaXUNyBMDfAvhVAIcBfIvkDjP7Tpn7FUlbxhjmqsbJ+hxZj7Go51K1jviU2k+e5JsA3GpmG4LvbwYAM7vNtb36yUu/uHrLt5YQIBYs1D3aGsk94i6Dr5f9CIk5M9XhN1xcP/my0zXjAP4n9P3h4L55JDeRnCA5MTk5WfLhiLi5qnFm5mxBgAcGq0InzHdBedZsIEtIpTp9nwxlZlsBbAU6I/k+H44MqSxVN1m2rWpCVpq2yd0TlGbeDpeyR/JtAGeEvj89uE9koGSpukm7bZUTslwXml2iJ6hBnzQm+ZUd5L8F4LUkV5M8DsC1AHaUvE+RzFxBsrWEaI0s7A2fpULHNyHrI184UHgVTLRax9XTHlh8ghr0SWOSX6npGjM7RvJ6ADsBjAC4y8wOlLlPkV74qnFc93W3TUpz+NI6zx+dwfNHZwAUWwUTrtbxLVIePUEN+qQxya/U6pqsVF0jdeELouHKG1/Fi8vYaAsnHL+00Lx4+CQ0tqwFM+CF6ZkFz+87xvGxUXxt88W59i/V6Wd1jUgjpUlzpM2TA8DU9EzhefGNa8fxtc0X485r1uCnM3OYmp5Z9PyDPmlM8lOQF+lBmjSHa1br2Ggr1fMXmRePOyGVMfNWBkvfSyhF6ijtEn/RWa2uNI9PUXnxpBNSkTNvZfBoJC/Sg17THK6R8/Jl7tF99ITRa28aX8nnElKlkkNAI3mRHvTSGyf8s0mj++gJI09vGteas0BnNqz62zSfqmtEBkBSOWaWKhhfo7Ubt+/DrOP/uypp6i+uukYjeZGKxAXypLx42np234j/tqvOxZxnQKea+GZTkBepQNp0i+9EkPZCb1wlTdrnkGbRhVeRCqSpq4/rI+O70HvR2SsWXIz1Tb76/tS0auKHlEbyIhVIk26JOxF0c+bhUf5FZ6/A3bvbCz4dEIArKXPa2Giui8VSXwryIhVIkyrxnQjaU9NYv2XXfGC+85o18y0JoicFAxYF+vBoPW1NfNb2w2pXPLiUrhGpQJpUiS83TsCZwvGdFAzINYM1a/thtSsebBrJi1QgTarEVc/uSr8kXUjNWxKZ1AYh7/ZSLQV5kYokpUqiJ4KxZa35lsRR7alpZx+cIi6kZm0/rHbFg03pGpEBEu0c6UN0OleGLV/WKqS5mC9tVNT9Ui0FeZEB5EqBdPkqaJYdt7SQ9EjWUkuVZg42pWtEBlBcqsPXiKSo9EjWUsteSjNVjVMdBXmRARR3URVA6TNXs7YfzrJ9nmZrkp3SNSIDKC4F0q/0SK+tjqOSZv8WtR/p0EheZAClSYFUme7I23snLK4aR6P84pXWapjkrQB+F8BkcNefmNn9cT+jVsMigylNq+M0i5snPRfgTkWpHXK8fi7kfaeZrQm+YgO8iAyuvL13wuLSTaq5L55y8iKSKE0tfNoAHbd4uGrui1d2Tv56kr8FYALAjWb2fMn7E5GUspQxulouuHrvpK368VXjpNlPL8c/zHKN5Ek+SHK/4+sKAB8D8BoAawAcAXCH5zk2kZwgOTE5OenaREQKlrWpWNzou6uIqp80++nl+IdZJWu8klwF4Itm9vq47XThVaQaWdaMzaKq0XVZx19XfVnjleRKMzsSfHslgP1l7UtEssmyZqxrUfBe16otii7QpldmTv4vSK5BZxb2UwDeW+K+RCSDNPlzV836Tf+6DzBgZs7m7+tHHbvWq02vtOoaM3uXmZ1rZr9oZpeHRvUi0mdp8ueuksiZWZsP8F3TM7O4cfu+SvPhaoqWnma8igyhNDNqs6Q+Zs0qHdFrvdr0KrnwmpYuvIr0TzT/fvSlY95FS3yG9cJnv/XlwquI1Icr/95aQrRGiJnZlweCrREuyMlHZb3wqVr38inIi4g7/z5nGBtt4YTjly6qrrlx+z7MOrIAWS58JjUj0wmgGAryIuIdgb8wPYO9H77E+VjSzNSkIJ3U66asbpTDdvJQ7xoRydwzJmlmapoZqXG17r4TwA3b9ubqMT+MM2U1kheRTD1juuImPvmC9I3b983/bFyte1xuP8+oPu7TQ1NTRBrJi0jqnjFp+YJ0t9Tynj1tZ6070QniS8jY53e1MM5zXOEFS8Kj/A9s24tVNV+hSiWUIlI4X2+ZrhESd1x9HoDO6Lo9NQ3Cv0i5CwF8b8vbCzmuuAVLwsZGW7j18nMGbnTfz0VDRGQIuUbpYeHJU1/bfDHGx0adAX4kZkTfSwuDXhYsCZuanqldDl9BXkQK103/xAXpcMrFF2DnzPDRa9YU1sKglwVL4o67DnThVURK0U1pRC/ohnWDe9xF2KJbGGRZsMSnTt0uFeRFpDTdYJo0eSqpuqeKFsbhk0nSNYI6dbtUkBeRUvlG9NEgDvQ2Wi+i7DH6HB+9Zg0A4CNfOLCof0/dul2qukZEKlFGDXq0NQLQCcJZyj+TnqMOtfNx1TUK8iJSW0UsA9iEpQRVQikijVTEMoBNX0pQOXkRqa0ilgFM+xx1SNu4KMiLSG1EA+1FZ6/A3bvbmXruRKXp25PUFnmQKV0jIrXg6i1z9+423nH+eK6eO2n69iS1RR5kGsmLSC34Au1Dj0/mvkCaVIdf57x9rpE8yXeSPEByjuS6yGM3kzxE8iDJDfkOU0SGXT8DbZp++/fsaWP9ll1YPWBdK/Oma/YDuArAw+E7Sb4OwLUAzgFwKYC/I+nvViQikiDrwiZFimtsBgz2YiS50jVm9hgAcHEToisAfNbMXgTwPZKHAFwA4D/z7E9EhlcvC5sUJTojd2xZC2bAB7btxe07D+LoS8cGdjGSsnLy4wC+Hvr+cHDfIiQ3AdgEAGeeeWZJhyMidZfU+qDsQNrN27sqbXzCi5H0qzInMciTfBDAqxwP3WJm9+Y9ADPbCmAr0Jnxmvf5RKS5fBdIqwykrgvAPqeNjSYuOVi2xJy8mb3VzF7v+IoL8G0AZ4S+Pz24T0SkcFWWOKa90Ju0GElVlTll1cnvAHAtyeNJrgbwWgDfLGlfIjLkqgykvgu9Y6OtTIuRVNWuOFdOnuSVAP4awAoA95Hca2YbzOwAye0AvgPgGIA/MLN0n29ERDIqor1BWr4LwL61X/t5wRhQF0oRaQBfu+B3nD+Ohx6fLPxibNaLvGVfFFarYRFpvHAgHVvWwoszszg6M7dgm7r1iU8rLsirrYGINMLiEse5RduEL8bWteFYVmpQJiKNklTi2J6axo3b99W24VhWCvIi0ihpKmpci4qn/dm6UbpGRBrFV2mT9merUtU1AY3kRaRRXM3E0qiyrLHKhmYayYtIo7h63PzkxWOYmp5ZtO0IiTmznkfSvY7Gq2x1oCAvIo0T7XHjq6PPuopUWJ5+OVXO0FW6RkQaL80Sf1nl6ZdTZasDjeRFZCgkLfGXVZ7ReJWtDhTkRUR6kNQvJy5fn9Qbv0hqayAi0oO4PD+Awq8BxIlra6CcvIhID3x5fgADNaNW6RoRkR75qngGaUatRvIiIgVJ6ptT5YzaLo3kRUQKEjdSH22N4KKzV2D9ll2VtjfWSF5EpCC+kfoIiXecP467d7craWUQpiAvIlIQV9+c0dYI7rj6PDz0+GRfLsYqyIuIFCRuZm2VrQzClJMXESmQb2ZtlYuNh+UayZN8J8kDJOdIrgvdv4rkNMm9wdfH8x+qiEh9+VI5Zbc3zjuS3w/gKgB/73jsSTNbk/P5RUQawdfKAECpFTe5gryZPQYAJIs5GhGRBktqgVzGguJlXnhdTXIPyX8n+cu+jUhuIjlBcmJycrLEwxERGSx52hWnlTiSJ/kggFc5HrrFzO71/NgRAGea2f+SPB/APSTPMbMfRTc0s60AtgKdBmXpD11EpN6qqLhJDPJm9tasT2pmLwJ4Mbi9m+STAH4egFpMiogEqqi4KSVdQ3IFyZHg9qsBvBbAd8vYl4hIXVVRcZO3hPJKkocBvAnAfSR3Bg/9CoBHSO4F8K8A3mdmz+U7VBGRZiljWcIoLRoiIlJzWjRERGRIKciLiDSYgryISIMpyIuINJiCvIhIgw1UdQ3JSQBP53iKUwD8sKDDKZKOKxsdVzY6rmyaeFw/Z2YrXA8MVJDPi+SEr4yon3Rc2ei4stFxZTNsx6V0jYhIgynIi4g0WNOC/NZ+H4CHjisbHVc2Oq5shuq4GpWTFxGRhZo2khcRkRAFeRGRBqtVkCf5TpIHSM6RXBd57GaSh0geJLnB8/OrSX4j2G4byeNKOs5tJPcGX08FLZdd2z1F8tFgu9Lbb5K8lWQ7dGyXeba7NHgdD5HcXMFx3U7ycZKPkPw8yTHPdqW/Xkm/O8njg7/voeC9tKqM43Ds9wySD5H8TvB/4P2ObS4k+ULo7/uhio4t9u/Cjr8KXrNHSL6hgmM6K/Q67CX5I5I3RLap5PUieRfJZ0nuD913MskHSD4R/Lvc87PXBds8QfK6ng7AzGrzBeAXAJwF4KsA1oXufx2AfQCOB7AawJMARhw/vx3AtcHtjwP4vQqO+Q4AH/I89hSAUyp8/W4F8EcJ24wEr9+rARwXvK6vK/m4LgGwNLj95wD+vB+vV5rfHcDvA/h4cPtaANsq+tutBPCG4PYrAfyX49guBPDFqt5Paf8uAC4D8CUABPBGAN+o+PhGAPwAnQlDlb9e6Kyv8QYA+0P3/QWAzcHtza73PICT0Vls6WQAy4Pby7Puv1YjeTN7zMxcK9xeAeCzZvaimX0PwCEAF4Q3IEkAF6OziAkAfArAxjKPN9jn1QA+U+Z+CnYBgENm9l0zewnAZ9F5fUtjZl82s2PBt18HcHqZ+4uR5ne/Ap33DtB5L70l+DuXysyOmNm3g9s/BvAYgOJWlijXFQD+yTq+DmCM5MoK9/8WAE+aWZ7Z9D0zs4cBRBdNCr+PfLFoA4AHzOw5M3sewAMALs26/1oF+RjjAP4n9P1hLP4P8DMApkLBxLVN0X4ZwDNm9oTncQPwZZK7SW4q+Vi6rg8+Mt/l+YiY5rUs07vRGfW5lP16pfnd57cJ3ksvoPPeqkyQIloL4BuOh99Ech/JL5E8p6JDSvq79Ps9dS38A61+vF4AcKqZHQlu/wDAqY5tCnndEhfyrhrJBwG8yvHQLWZ2b9XH45PyOH8D8aP4N5tZm+TPAniA5OPBWb+U4wLwMQB/hs5/yj9DJ5X07jz7K+K4uq8XyVsAHAPwac/TFP561Q3JEwHcDeAGM/tR5OFvo5OS+L/gess96KyvXLaB/bsE190uB3Cz4+F+vV4LmJmRLK2WfeCCvJm9tYcfawM4I/T96cF9Yf+LzsfEpcEIzLVNaknHSXIpgKsAnB/zHO3g32dJfh6ddEGu/xxpXz+S/wDgi46H0ryWhR8Xyd8G8GsA3mJBQtLxHIW/XhFpfvfuNoeDv/FJ6Ly3SkeyhU6A/7SZfS76eDjom9n9JP+O5ClmVmozrhR/l1LeUym9DcC3zeyZ6AP9er0Cz5BcaWZHgtTVs45t2uhcN+g6HZ3rkZk0JV2zA8C1QeXDanTOxt8MbxAEjocA/Hpw13UAyvxk8FYAj5vZYdeDJE8g+crubXQuPu53bVuUSB70Ss/+vgXgtexUIh2HzkfdHSUf16UA/hjA5WZ21LNNFa9Xmt99BzrvHaDzXtrlOykVKcj7fwLAY2b2l55tXtW9PkDyAnT+f5d6Akr5d9kB4LeCKps3AnghlKoom/fTdD9er5Dw+8gXi3YCuITk8iC1eklwXzZlX1ku8gudwHQYwIsAngGwM/TYLehURhwE8LbQ/fcDOC24/Wp0gv8hAP8C4PgSj/WTAN4Xue80APeHjmVf8HUAnbRF2a/fPwN4FMAjwZtsZfS4gu8vQ6d648mKjusQOrnHvcHXx6PHVdXr5frdAfwpOicgAHhF8N45FLyXXl326xPs983opNkeCSYaCkYAAACPSURBVL1OlwF4X/d9BuD64LXZh84F7F+q4Licf5fIcRHA3wav6aMIVcaVfGwnoBO0TwrdV/nrhc5J5giAmSB+vQed6zhfAfAEgAcBnBxsuw7AP4Z+9t3Be+0QgN/pZf9qayAi0mBNSdeIiIiDgryISIMpyIuINJiCvIhIgynIi4g0mIK8iEiDKciLiDTY/wP3TP2CPf08HwAAAABJRU5ErkJggg==\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ne9UvrY7cIol"
      },
      "source": [
        "Let us set up initial position in parameter space of (m, b). This initial point corresponds to a 45 degree line fot which is clearly bad."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3dJN3zg-APQx",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 267
        },
        "outputId": "6cdccfa0-abc2-4249-9557-9e8e30e69e4b"
      },
      "source": [
        "m_model = 1.0\n",
        "b_model = 0.0\n",
        "\n",
        "y_pred = y(x, m_model, b_model)\n",
        "plt.plot(x, y_true, 'o')\n",
        "plt.plot(x, y_pred, 'o')\n",
        "plt.show()"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXkAAAD6CAYAAABEUDf/AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3df5RcdZnn8feTToMNuulkySA0ZBLnuDACmkiWHQdnD4gCOiPEuAK758wwq8fo7rqOHpYhDOdA1D8SZBlcd2d04oxHneMqURBRdCMKMx45o2NCfgACCwouaRHiQKOSHuh0P/tH3ercrr731r11f1TdW5/XOTnprqque1Ndee63nu/zfb7m7oiISDMt6fcJiIhIeRTkRUQaTEFeRKTBFORFRBpMQV5EpMEU5EVEGix3kDezk83sbjP7kZk9YGZ/Ety+wszuNLNHgr+X5z9dERHJwvLWyZvZCcAJ7n6vmb0M2A1sAP4YeMbdt5nZZmC5u1+V9FzHHXecr169Otf5iIgMm927d//C3VdG3bc075O7+5PAk8HXvzKzB4EJ4GLgnOBhnwX+DkgM8qtXr2bXrl15T0lEZKiY2U/j7is0J29mq4F1wA+A44MLAMDPgeNjfmaTme0ys10HDx4s8nRERIZeYUHezF4K3AJ8wN1/Gb7PWzmhyLyQu2939/Xuvn7lyshPGyIi0qNCgryZjdIK8J9391uDm58K8vXtvP3TRRxLRETSK6K6xoC/AR509z8P3XU7cHnw9eXAV/MeS0REssk98QqcDfwhcJ+Z7Q1u+zNgG7DDzN4F/BS4pIBjiYhIBkVU13wPsJi7z8v7/GnctmeSG3Y+zM+mpjlxfIwrLziFDesmqji0iMhAK2Ik31e37Znk6lvvY3pmFoDJqWmuvvU+gNhAr4uCiAyL2gf5G3Y+PB/g26ZnZrlh58MLAnc7sE9OTWMcKfVJc1EQEamr2veu+dnUdNfb26P9yeC2zlrO9kVBRKRpah/kTxwf63p71Gi/U9zFQkSkzmof5K+84BTGRkcW3DY2OsKVF5wy/32aAB53sRARqbPa5+TbefSkidQTx8fmUzVRwhcFTcqKSJPUPshDK9AnBeIrLzhlQQUOMD/5OhEK5L1U6oiIDLJGBPlu0oz22/dHVepcsWPfgucREamLoQjy0H20D/G5+1l3jehFpJZqP/FapKTJV5VZikgdKciHRFXqhKnMUkTqRkE+ZMO6CbZuPIMRi27FozJLEakbBfkOG9ZNcOMlr+laey8iUgdDM/GaRdpqHBGRQacgHyNNNY6IyKBrfJDXClYRGWaNDvJawSoiw67RQT5tr/lu9GlAROqq0UE+Ta/5bsr6NKALh4hUodEllGl6zXeT9GmgV+FNTJwjF47b9kz2/JwiIlEKCfJm9mkze9rM7g/dtsXMJs1sb/DnLUUcK4s0vea7KeLTQKcyLhwiIlGKGsl/Brgw4vab3H1t8OcbBR0rtfYK1onxMYxWW+GtG8/IlBaJG/UvMet55F3GhUNEJEohOXl3/66ZrS7iuYqWt949qhc95OtMGbeJybKxUc7edpfy9CJSmLJz8u8zs/1BOmd5yccqRVI/m15TLFFppNElxvMvHlaeXkQKVWaQ/wTwW8Ba4EngxqgHmdkmM9tlZrsOHjyY64C37Znk7G13sWbzHZy97a7CAuSGdRPMuUfe10uKJSqN9NKXLGVmduExlKcXkbzMY4JX5idqpWu+7u6nZ7kvbP369b5r166ejt9Z6tg2PjbKlotOy532OHvbXZEplonxMe7Z/IZczw2wZvMdRP0mDHhs2+/nfn4RaS4z2+3u66PuK20kb2YnhL59G3B/3GOLEFWxAjA1PVNI2qOISp0o7U8fcZdatTcWkTwKmXg1sy8A5wDHmdkB4DrgHDNbS2u/7MeB9xRxrDhJaZNeVrl2KqMzZdynjza1NxaRvApL1xQhT7omLp3SVmTaI+tq1bjHJ53zhKprRCSlpHRNY9oaxJU6thWV9sja5iDp8XGfPgzm8/xqfyAieTSmrUG7YmX5MaOL7isy7ZF1tWrS47u1XVD7AxHJqzFBHlqBfs+15/OxS9fmWuWaJOtq1aTbu03mqv2BiOTVmHRNWNpVru1UyOTUNCNmzLp3zYXHrVZNGpXHPb7bZK7aH4hIXo0M8ml05spngwnobjn2qNx/Ujqo2+OTLkhZLygiIp0ala7JIq6uHpJTIlmanrU/KUzPzM63RciSPuqlNr+sVb8iUk9DO5LvlvJIuj88+m4H8g/evHdBuiXqk0I7QKedH8ham6/tDkWkU2Pq5LPqVlefpl1B3GKmJQZzMS9rUW0QosT9m0bMmHNXCaZIQ/WlrcGgi0qFtKUtuYxL+cQFeCh30jTuuWfdVYIpMqSGNsiHc+tATznzXgJ2mZOmaZ5bJZgiw2Voc/KQf0ORuOqXOGX3oum26rdNJZgiw2NoR/JFSEr5dCpqUVZS9Uxn5U/URiegEkyRYTLUI/m82gF7y+0PMDU9E/mYsdGRwoL7h772AM8eOnKcqOqZzsqftDX96pEj0kwayee0Yd0Ee6870koBesvvJ2kH63CAbyuipl89ckSaa2hLKOuk7DbKZe96JSLlGopWw03WbaJ02dgoZ2+7q+dUi3rkiDSX0jU1kDRROrrEeP7Fw7lSLd1aHotIfSnI10BcFc/42CgvfclSZmYXptyy1sKXtX+tiPSf0jU1kNTDZs3mOyJ/JkuqJe75gVxpIBHpPwX5mohbuFVUO+LO51ezM5FmULqm5spKtWhXKpFmKGQkb2afBv4AeNrdTw9uWwHcDKwGHgcucfdnizhe02VZmJS1HXFaqrgRqcj+HfCdD8NzB2DZSXDetfDqSwp7+qLSNZ8B/hfwudBtm4HvuPs2M9scfH9VQcerrW4BvJc0Sd4ePFG0K5VIyfbvgG9eBdPPHLntuSfga+9vfV1QoC8kXePu3wWe6bj5YuCzwdefBTYUcaw6S7OydFDSJKq4ESnB/h1w0+mwZRncumlhgG+bmW6N7AtS5sTr8e7+ZPD1z4Hjox5kZpuATQCrVq0q8XT6LymAD9rm3WWlgUSGRjgNM7YcDr8AM8+HHpDQbeC5A4WdRiXVNe7uZhb5L3L37cB2aLU1qOJ8+iVNAK8iTZI25x+VBlIjM5EuotIwUSP2JMtOKux0yqyuecrMTgAI/n66xGPVQpqVpWWnSfI0I1MjM5EE+3fA9Wvg1ndnD+pho2OtydeClBnkbwcuD76+HPhqiceqhTQBPG3nyF7lyfn38rNJ/e9Fai9Njj2LsRXw1o8PXnWNmX0BOAc4zswOANcB24AdZvYu4KdAcWddU2nz3GmrZZJSJ3H35cn5d/vZzmOee+pKbtk9qQVV0jxRKZmkHHssa/3cspMLL52cP4JaDddT3IYgWzeeARB73w07H+65rXBcS+IRM2bd22/XeZ3fZzmWyMCZn0h9gvh3dwZjK+DN1xcS2NVquIG6pU7i7ovaBzZtzj9uD9nZYKDQ+ZaP+y+Q9KlBE7sycIoYtY8eC0uPhulnS1nwlERBvqZ6Sbv8bGo6V2lk588uCUbwWcVNQKtfjgyUyOCeUYGj9V4pyNdUt1LLqPs6Nxe56dK1mYNneL4grgNmkqRPDWnWEYiUKndKpvwce1YK8jXVLe3SeV97c5H2huNFjJLjLjRxJrp8ahiUhWAyhIpIyQzAqD2KgnxNpUm7hO879OLhRRuB5x0lx+XooxgsmmztzL8vGxudvwiFqV+OlKKBo/YoCvI1llRq2XlfEZuLRB0Dul9MYHGgjsq/j44Yo0uMmbkj/9nUL0cK1+BRexQF+SFRVruEbpuNQHSgjsq/z8w6y48Z5Zijlqq6RorRtX9MRjUK7m0K8kMiT+lkFmmrd+I+QUwdmmHPteenPp5KLiVSEf1jgLqkZJIoyA+JKvdxTbNit4hPFiq5lAUGeLFSP2nF6xCLSq2MLjFe+pKlTB2aKXVknLRiN+3x4lbgakXtkCminh2o86hdK14lUmRefM7nJ07LHBkX0a9eJZdDrrDgTmNG7VEU5IdYmmBY5mKkvNsWaovCIaSUTGYK8kMs7WKmqItBPyY828ecnJqObYqmksuGqnn/mH5SkB9iaRczpalxL3vCs/OY4aZo7UDfbUWt1MyQLFYqm4L8EOvMiy8bG+X5Fw8zM5u8GKkfPWaijtnWDvBxk60qs6yZIVusVDYF+SEXtZip1xr3Mic8uz335NR0ZClo3KeOXT99hrsfOqjAP0ga0vVx0CjIywJV1bhn1W3+wDjSeTOcPor71PH57/+/+bGh6uv7SCmZ0inIS2ZVrZ6FhZOtcSEg6vZ2+ijuE0Dc4xXkK6KUTGUU5KWrqBROeyvBMtMdnamW8CRru7pmImGE3z63tO2QVV9fkvBo3UbAZ9GovToK8pIoLqe9deMZpa8qjUq1RE2yxq18bV98Oj91xIUX1dcXLGq07u3fg0btVSk9yJvZ48CvgFngcNzSWxlM/dytKe0Eb1L6KGpl7bmnruSW3ZOVpJuGjhYrDZyqRvLnuvsvKjqWFKifrQPSTvB2a5EQNZm8/jdXqKyySEXk2AGlZIqndI0kKqqSppda9SwTvFlbJORtqSBo1F4TVQR5B75lZg78lbtvD99pZpuATQCrVq2q4HQkiyIqaXpdIVtEEzMpgUbttVJ6q2Ezm3D3STP7DeBO4L+6+3ejHqtWw4Mp74rRuInRETNuvOQ1Ctp1UcRiJVsCPqfAXrC+thp298ng76fN7CvAWUBkkJfBlDe1EZe/n3WvzSKkoW2NoMVKtVdqkDezY4El7v6r4OvzgQ+XeUwZPEm16nVYhDSUO1BpsVJjlD2SPx74ipm1j/W/3f3/lHxMGTDdul1WuQiplxF5P8tIK6VReyOVGuTd/SfAa8o8hgy+diC8Yse++RbBYePHjBa+z2yUXkfkcRehyalpbtszCdR8clij9kZTCaVUoh30Fu0pO2L8+p8PL9pysIwukb2OyJPSTVd+aR8Y8+2Za5XKUdfHoaAgL5WJKol8/oXDTE3PLHhcWV0ie13YlZRumplbPOIdyFSO+scMLQV5qVRnpc6azXdEPq7XLpFJOfdeF3a1f/4DN+9NfFzYQDQ7i8uxq3/MUFnS7xOQ4ZZl5Wy3wNnOuU9OTeMc+QTQzptfecEpjI2OLPiZtAu7NqybYCLDufa12dn+HXD9Grj13UGAh94XK9EatW/8FFz1mAJ8DSnIS19FBV6LeWy3wJmUc4dWoN668QwmxscwWt0st248I3VaJepcR5cYoyMLz7i9gcnZ2+6av8BUIhzc8+TZoTVq37gdtjwHH7xfwb3GlK6RviqyS2SanHvUdodpK3vi2iy0b+vc2KSSSVj1j5EuSm9rkIXaGkhbL/Xsce0T4jb57iyphNbF5O1nTvRU2ZP1+LkUURkDaCK1Gfra1kCkF720UsjaTK3o/V9Lb8tc1GKldnWNAvtQUJCXxsjatbLo/V/jqneWmHHbnsneUzZarCQ5KMhLo2T5BFD0/q9x9fQ9N2LTYiUpgIK8DK2i939Nat+QeoGU+sdIwRTkZWiVsf/rhnUTfDBm0VTipwGlZKQkCvIy1MrY/7Xrylq1GJAKKciLdMi7SUpclc/HXvUIXP+uhaN1tRiQkinIixQsnAZa/8s7ufqoL3E8B7F7tVhJqqcgL43Vzy37Nozcwwa7Co7SZtfSXwry0kh927KvsJWoaNQuhVCQl0aqdMu+gvrHzHnrp21co3YpjoK8NFLpLQYgd9mjO8xhLMGZ9OP46OFL2P0v3sQ9Hyy4z40MNQV5aaReNwjpqqDFSofGTuDa59/Ol1/83fl7xkZH2JqyHl8krdKDvJldCPwPYAT4a3ffVvYxRbI2K+uq4MVKxwCv3zPJPxQ4MdzPiWYZXKUGeTMbAf4CeBNwAPihmd3u7j8q87giWZuVQXSQnHji67zy3o+wzH+Fxe1m0k3MBGreevywvk00y8ArtZ+8mb0O2OLuFwTfXw3g7lujHq9+8tIv4SB50ZLv8adLd3Ci/QKAJZmDe/Vlj3G97EfMmHPXyL7h+tlPfgJ4IvT9AeDfhB9gZpuATQCrVq0q+XREot2w82HeNPv3XHfU51hhvy581F62uAnldqM0jeyHV98nXt19O7AdWiP5Pp+ODIuO/jHfm5vFR4sftVeVJ0/TNjmuhFS5/GYrO8hPAieHvj8puE2kP6ImUH0Ws/gNxGN1GbVXmSeP62XfqXPEr1x+85Ud5H8IvNLM1tAK7pcB/6HkY4osVOBm1+7w4lHLOPqt/71rSiZuQdaHvvZA4SPnzonmJWaLetrD4hLSSheNSV+UGuTd/bCZvQ/YSauE8tPu/kCZxxSZV0TZIzAHmMNTtpInzrySf33Re4DuaY64PPmzh2Z49tAMUOzIOVytE7dJeWcJaSWLxqSvSq2uyUrVNZJbgaN2IDYlExdEt248Yz7QxlW8RBkfG+XYo5cWOroPX4TGjxnFHZ6bnlnw/HHnODE+xj2btfK2LpKqa5ZUfTIipdi/A65fA7e+Owjw0HvXR1oTqRs/BVc9FpmWSUpztF15wSmMjY6kOurU9AyTU9M4R0b3t+3JN321Yd0E92x+AzddupZ/npljanpm0fNHnWOuRWMycPpeXSOSSxFdH20J+FymuvY0aY6oBVnPv3CYqemZrs9fZF486YLUHq2ruqa5FOSlfgZgs+u0vXE6V7VGpXniFJUX73ZBKnLlrQweBXmpjwHa7LrX3jhRo/tDLx6en4gN67xg9FrPHndBWmLGbXsmFeAbTkFeBtsAjNqj9NIbJ/yz3Ub3nReMPPXscTX0s+6qiR8Cqq6RwVRErr1GOyt1G6VnqYKJei6AK3bsi6ydVyVN/fWzd41INg0O7kmBvFtePG09e9yIf+vGM5iLGdCpJr7ZFOSlfzr6x+CzDFJKpkhp0y1xF4K0E71JlTSlbaQiA01BXqoVl2P3dmDqz0Rq2dK0D0i6EMRN9J576krO3nbX/EUhbvHVz6amuenStcVupCK1oCAv1SioxUAdRu1R0qRbstazn3vqSm7ZPbngohD3OejE8bFck8VSXwryUq4icuxtNRm1R0mTKom7EExOTS8Yrd906dr5lgSdFwVnccIrPFpPWxOftVxT7YoHl4K8FK+i/jF1kqauPu5CYDB/eziFE3dRcFoVM70G3KzlmmpXPNgU5KU4Q56SSZImVRJ1IYi6RHabSM1bEpm1/bDaFQ82BXnJp6jFSu3qmgYF9k7dUiWdF4LxY0YjV8JCa7Q8Pja66PYiJlKzth9Wu+LBpiAvvRmgFgNN0r4QtFMgcQwWNTpbfswo1731tNyj56yllirNHGxqNSzZhFv69jqZOrYisY2vRKdA2uI+Lx1z1NJC0iNZ2w+rXfFg00heuhvQ/jFNlpTqiHv1i0qPZC217KU0U9U41VGQl3hKyfRN0qQqUHp6JGv74SyPVzVOtRTkZSGN2gdCt5LLfqxcLWr03a0aR6P8YinID7vEoK5Re7+kSYFUGQjz9t4JS6rG0Si/eKW1GjazLcC7gYPBTX/m7t9I+hm1Gq6IFitJRmlaHafZ3Lzbc0F0KkrtkJP1cyPvm9x9bfAnMcBLBSre7FqaI2/vnbCkahzV3BdP6Zqm06hdCpCn907n7UmpqBt2Pqya+4KVHeTfZ2Z/BOwCrnD3Z0s+nrSpxYB0kWWCM0/vnagAHVeNk2XvXE3QppMryJvZt4GXR9x1DfAJ4CO0IstHgBuBd0Y8xyZgE8CqVavynI6Auj5KKlknOHvtvZO16idtzb0maNOrZI9XM1sNfN3dT096nCZee6T+MZJRlj1js6hqdF3W+ddVX/Z4NbMT3P3J4Nu3AfeXdayhpcVK0qMse8ZGbQre6161RdEEbXpl5uQ/amZraUWdx4H3lHis4aHFSlKANPnzqJTIlV/eBw4zcz5/Wz/SJGqKll5pQd7d/7Cs5x5KGrVLgdLkz6NKImdmF7/npmdmuWLHPqC6QF9E/n9YqIRy0BUxkargLh3STHBmSX3Mulc6otd+telVMvGa1tBPvIZTMe1JUKVkpCKd+fdDLx6O3bQkzrBOfPZbXyZeJaW4HLu3P4YqJSPli8q/jy4xRkdsQYpmdMQW5OQ7ZZ34VK17+RTk+0WLlWSARObf55zxsVGOPXrpouqaK3bsYzYiC5Bl4rNbrbsuAMVQkK+aFivJAIobgT83PcPe686PvK/bxGe3IN2t101Zi52G7eKhIF8F9Y+RAZe1JLHbxGeaFalJte5xF4AP3LyXG3Y+3HNgHsaVsgryZVJKRmqil5LEpIVPcUE6XGqZdGFJyu3nCczDuGGJqmuKphYDUlNFBrg1m++Ifee3e8zD4pRP+3/MiFlkzj+sl0qeuPMy4KZL18aez8SAB/yk6hoF+aKonl1kXlxvmbYRM2685DUA8+2Fsw6JDHhs2+8Xcl5JG5aEjY+NsuWi0wYu2Pdz05DmC2/E0WuAH1uhzTekUaI2BgkLL566Z/MbmBgfiwzwI2axz9FLC4NeNiwJm5qe4epb7+O2PZOZj90vCvK92L8DbjodtiyDWzf1ENy1s5I024Z1E2zdeEZikA5X0sQF2Dl3Pnbp2tjA3Ot5TYyPYbRG8O3tCdNeNKJ2uxpkmnjNQv1jRFJrpzQ689xh7eCeNAlbdAuDLBuWxKlTt0sF+TRy5dtVGSPDqx1Muy2e6lbdU0UL4/DFpNscQZ26XSrIdyqsfwwatYsQP6LvDOLQ22i9iKqgzuf42KVrAfjQ1x5Y1L+nbt0uVV0DWqwkUoEyatA7FzfBkRLNtM/d7TnqUDuvEso4hbUYUEpGpB+K2AawCVsJqgtlJ/WPEWmEIrYBbPpWgsMT5JWSEWmcIrYBTPscdUjbRGl+kFf/GJHG6Ay05566klt2T+baBjBN3546NzZrZpBX/xiRxokKtLfsnuTtZ05w90MHex5hp6ns6dbYbJA1K8hrsZJIY8UF2rsfOph7grRbHX6d8/a5gryZvQPYAvw2cJa77wrddzXwLmAWeL+778xzrERqDibSeP0MtGny9oOas8/bu+Z+YCPw3fCNZvYq4DLgNOBC4C/NLL5bUR77d8DX3q/+MSINFzeZWsXq06TGZnAklTQ5NY1zJGc/CI3Mco3k3f1BAFvchOhi4Ivu/gLwmJk9CpwF/EOe40X6zodhJuOVXKN2kdrpZWOTonTm7cePGcUdPhjsVHXoxcMDuxlJWTn5CeD7oe8PBLctYmabgE0Aq1atyn6k5w6keJAqY0TqLs2Wg2UG0nbePmoCOM7Ppqb7XpnTNcib2beBl0fcdY27fzXvCbj7dmA7tFa8Zn6CZScFVTQxNGoXaYy4CdIqA2nUBHCcE8fH+l6Z0zUn7+5vdPfTI/4kBfhJ4OTQ9ycFtxXvvGthNCInp404RIZGUiAtWtqJ3m6bkVRVmVPWpiG3A5eZ2dFmtgZ4JfCPpRzp1ZfAWz/eSsVgmkgVGUJVBtK4id7xsdFMm5FU1a44bwnl24D/CawE7jCzve5+gbs/YGY7gB8Bh4H/4u7pPt/04tWXKKCLDLEi2hukFTcBHLf3az8njGHYu1CKSCPEtQvOuxo26XhZJnnLnhRWq2ERabxwIB0/ZpQXZmY5NDO34DF16xOflloNi0jjLS5xnFv0mPBkbF0bjmVV1sSriEhfdCtxnJya5ood+yqrxuk3BXkRaZQ0FTVRm4qn/dm6UbpGRBolrtIm7c9Wpao5AY3kRaRRopqJpVFlWWOVDc00kheRRonqcfP8C4eZmp5Z9NgRM+bcex5J9zoar7LVgYK8iDROZ4+buDr6djllL/L0y6lyha7SNSLSeBvWTbB14xmRbQd6ladfTpWtDjSSF5Gh0G2Lv6zyjMarbHWgIC8i0oNu/XKS8vVpNg8vitoaiIj0ICnPDxQ+B5Akqa2BcvIiIj2Iy/MDA7WiVukaEZEexVXxDNKKWo3kRUQK0q1vTpUrats0khcRKUjSSH1sdIRzT13J2dvuqrS9sUbyIiIFiRupj5jx9jMnuGX3ZCWtDMIU5EVEChLVN2dsdIQbL3kNdz90sC+TsQryIiIFSVpZW2UrgzDl5EVEChS3srbKzcbDco3kzewdZvaAmc2Z2frQ7avNbNrM9gZ/Ppn/VEVE6isulVN2e+O8I/n7gY3AX0Xc92N3X5vz+UVEGiGulQFQasVNriDv7g8CmFkxZyMi0mDdWiCXsaF4mROva8xsj5n9vZn9XtyDzGyTme0ys10HDx4s8XRERAZLnnbFaXUdyZvZt4GXR9x1jbt/NebHngRWufs/mdmZwG1mdpq7/7Lzge6+HdgOrQZl6U9dRKTeqqi46Rrk3f2NWZ/U3V8AXgi+3m1mPwb+FaAWkyIigSoqbkpJ15jZSjMbCb5+BfBK4CdlHEtEpK6qqLjJW0L5NjM7ALwOuMPMdgZ3/Vtgv5ntBb4MvNfdn8l3qiIizVLGtoSdtGmIiEjNadMQEZEhpSAvItJgCvIiIg2mIC8i0mAK8iIiDTZQ1TVmdhD4aY6nOA74RUGnUySdVzY6r2x0Xtk08bx+091XRt0xUEE+LzPbFVdG1E86r2x0XtnovLIZtvNSukZEpMEU5EVEGqxpQX57v08ghs4rG51XNjqvbIbqvBqVkxcRkYWaNpIXEZEQBXkRkQarVZA3s3eY2QNmNmdm6zvuu9rMHjWzh83sgpifX2NmPwged7OZHVXSed5sZnuDP48HLZejHve4md0XPK709ptmtsXMJkPn9paYx10YvI6PmtnmCs7rBjN7yMz2m9lXzGw85nGlv17d/u1mdnTw+300eC+tLuM8Io57spndbWY/Cv4P/EnEY84xs+dCv99rKzq3xN+LtXw8eM32m9lrKzinU0Kvw14z+6WZfaDjMZW8Xmb2aTN72szuD922wszuNLNHgr+Xx/zs5cFjHjGzy3s6AXevzR/gt4FTgL8D1odufxWwDzgaWAP8GBiJ+PkdwGXB158E/lMF53wjcG3MfY8Dx1X4+m0B/luXx4wEr98rgKOC1/VVJZ/X+cDS4Ovrgev78Xql+bcD/xn4ZPD1ZcDNFf3uTgBeG3z9MuD/RpzbOcDXq3o/pf29AC8id1UAAARESURBVG8BvgkY8DvADyo+vxHg57QWDFX+etHaX+O1wP2h2z4KbA6+3hz1ngdW0NpsaQWwPPh6edbj12ok7+4PunvUDrcXA1909xfc/THgUeCs8APMzIA30NrEBOCzwIYyzzc45iXAF8o8TsHOAh5195+4+4vAF2m9vqVx92+5++Hg2+8DJ5V5vARp/u0X03rvQOu9dF7wey6Vuz/p7vcGX/8KeBAobmeJcl0MfM5bvg+Mm9kJFR7/PODH7p5nNX3P3P27QOemSeH3UVwsugC4092fcfdngTuBC7Mev1ZBPsEE8ETo+wMs/g/wL4GpUDCJekzRfg94yt0fibnfgW+Z2W4z21TyubS9L/jI/OmYj4hpXssyvZPWqC9K2a9Xmn/7/GOC99JztN5blQlSROuAH0Tc/Toz22dm3zSz0yo6pW6/l36/py4jfqDVj9cL4Hh3fzL4+ufA8RGPKeR167qRd9XM7NvAyyPuusbdv1r1+cRJeZ7/nuRR/OvdfdLMfgO408weCq76pZwX8AngI7T+U36EVirpnXmOV8R5tV8vM7sGOAx8PuZpCn+96sbMXgrcAnzA3X/Zcfe9tFISvw7mW26jtb9y2Qb29xLMu10EXB1xd79erwXc3c2stFr2gQvy7v7GHn5sEjg59P1JwW1h/0TrY+LSYAQW9ZjUup2nmS0FNgJnJjzHZPD302b2FVrpglz/OdK+fmb2KeDrEXeleS0LPy8z+2PgD4DzPEhIRjxH4a9XhzT/9vZjDgS/42W03lulM7NRWgH+8+5+a+f94aDv7t8ws780s+PcvdRmXCl+L6W8p1J6M3Cvuz/VeUe/Xq/AU2Z2grs/GaSuno54zCSteYO2k2jNR2bSlHTN7cBlQeXDGlpX438MPyAIHHcD/y646XKgzE8GbwQecvcDUXea2bFm9rL217QmH++PemxROvKgb4s53g+BV1qrEukoWh91by/5vC4E/hS4yN0PxTymitcrzb/9dlrvHWi9l+6KuygVKcj7/w3woLv/ecxjXt6eHzCzs2j9/y71ApTy93I78EdBlc3vAM+FUhVli/003Y/XKyT8PoqLRTuB881seZBaPT+4LZuyZ5aL/EMrMB0AXgCeAnaG7ruGVmXEw8CbQ7d/Azgx+PoVtIL/o8CXgKNLPNfPAO/tuO1E4Buhc9kX/HmAVtqi7Nfvb4H7gP3Bm+yEzvMKvn8LreqNH1d0Xo/Syj3uDf58svO8qnq9ov7twIdpXYAAXhK8dx4N3kuvKPv1CY77elpptv2h1+ktwHvb7zPgfcFrs4/WBPbvVnBekb+XjvMy4C+C1/Q+QpVxJZ/bsbSC9rLQbZW/XrQuMk8CM0H8eheteZzvAI8A3wZWBI9dD/x16GffGbzXHgX+Yy/HV1sDEZEGa0q6RkREIijIi4g0mIK8iEiDKciLiDSYgryISIMpyIuINJiCvIhIg/1/7o4hA+BrKxkAAAAASUVORK5CYII=\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A5sJH5oAHrSq"
      },
      "source": [
        "## 2: Set up a loss function and get its gradient\n",
        "\n",
        "We'll define our loss function in terms of our parameters, so that we can calulate the gradient with respect to them."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NFZinhp0Ak9T"
      },
      "source": [
        "from jax import grad, jit\n",
        "\n",
        "def sqr_loss(y_true, y_pred):\n",
        "  return jnp.mean((y_true-y_pred)**2)\n",
        "\n",
        "def loss_wrt_params(x, m, b, y_true):\n",
        "  y_pred = y(x, m, b)\n",
        "  return sqr_loss(y_true, y_pred)\n",
        "\n",
        "# note: by dft grad calculates gradients w.r.t the first arg\n",
        "#  but in this case we want it w.r.t. m & b, i.e. args 1 & 2\n",
        "gradient_wrt_loss = jit(grad(loss_wrt_params, argnums=(1, 2)))\n"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice the function composition. We jitted on the grad on the loss function. The `argnums=(1, 2)` describe which arguments are the parameters in the loss function `loss_wrt_params`. Here, the second and third arguments.\n",
        "\n",
        "We can calculate the gradient at the true model parameters as well as the initial starting point in parameter space."
      ],
      "metadata": {
        "id": "LrjSxB3XSmVQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "m_grad, b_grad = gradient_wrt_loss(x, m_true, b_true, y_true)\n",
        "print(\"for true: m_grad\", m_grad , \"b_grad\", b_grad)\n",
        "\n",
        "m_grad, b_grad = gradient_wrt_loss(x, m_model, b_model, y_true)\n",
        "print(\"for model params: m_grad\", m_grad , \"b_grad\", b_grad)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5AC0O8OKSnH5",
        "outputId": "e0fde4ec-6122-414a-e958-82029503dbea"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "for true: m_grad 0.8275 b_grad -0.06965329\n",
            "for model params: m_grad 170.11884 b_grad -4.069654\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yaz3COkbd3s9"
      },
      "source": [
        "Finally this gradient can be used in a for loop to implement gradient descent!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aSN8twa5DB5k",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cee81153-f8a9-4bc6-81ba-06e621fa9b2a"
      },
      "source": [
        "m, b = 1.0, 0.0\n",
        "learning_rate = 0.01\n",
        "\n",
        "print(\"true\", m_true, b_true)\n",
        "for _ in range(30):\n",
        "  m_grad, b_grad = gradient_wrt_loss(x, m_model, b_model, y_true)\n",
        "  m_model -= learning_rate * m_grad \n",
        "  b_model -= learning_rate * b_grad \n",
        "  print(\"grads\", m_grad, b_grad , \"=> model\", m_model, b_model)"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "true -1.5 2.0\n",
            "grads 170.11884 -4.069654 => model -0.7011883 0.04069654\n",
            "grads 54.920258 -3.988261 => model -1.2503909 0.08057915\n",
            "grads 17.730167 -3.908495 => model -1.4276925 0.119664095\n",
            "grads 5.723911 -3.830325 => model -1.4849316 0.15796734\n",
            "grads 1.8478823 -3.7537184 => model -1.5034105 0.19550453\n",
            "grads 0.5965558 -3.678644 => model -1.509376 0.23229097\n",
            "grads 0.19258726 -3.6050706 => model -1.5113019 0.26834166\n",
            "grads 0.062177062 -3.5329704 => model -1.5119237 0.30367136\n",
            "grads 0.020069182 -3.4623103 => model -1.5121244 0.33829445\n",
            "grads 0.006476879 -3.393065 => model -1.5121891 0.3722251\n",
            "grads 0.0020920038 -3.325204 => model -1.51221 0.40547714\n",
            "grads 0.00067955256 -3.2586987 => model -1.5122168 0.43806413\n",
            "grads 0.00022101402 -3.1935244 => model -1.5122191 0.46999937\n",
            "grads 6.657839e-05 -3.1296546 => model -1.5122198 0.5012959\n",
            "grads 1.8715858e-05 -3.0670617 => model -1.51222 0.53196657\n",
            "grads 2.3245811e-06 -3.0057197 => model -1.51222 0.56202376\n",
            "grads 1.7881393e-06 -2.9456053 => model -1.51222 0.59147984\n",
            "grads 1.4901161e-06 -2.8866937 => model -1.51222 0.6203468\n",
            "grads 2.2053719e-06 -2.8289592 => model -1.51222 0.6486364\n",
            "grads 2.1457672e-06 -2.7723799 => model -1.51222 0.6763602\n",
            "grads 1.7285347e-06 -2.7169323 => model -1.51222 0.70352954\n",
            "grads 2.1457672e-06 -2.662594 => model -1.51222 0.73015547\n",
            "grads 2.1457672e-06 -2.609342 => model -1.51222 0.7562489\n",
            "grads 2.5331974e-06 -2.5571558 => model -1.51222 0.7818205\n",
            "grads 2.503395e-06 -2.5060122 => model -1.51222 0.8068806\n",
            "grads 2.0861626e-06 -2.4558916 => model -1.51222 0.8314395\n",
            "grads 3.0994415e-06 -2.4067743 => model -1.51222 0.85550725\n",
            "grads 2.6524067e-06 -2.3586392 => model -1.51222 0.87909365\n",
            "grads 2.592802e-06 -2.3114662 => model -1.51222 0.9022083\n",
            "grads 2.2947788e-06 -2.2652369 => model -1.51222 0.9248607\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ar4lT38IBGzj"
      },
      "source": [
        "## 2: pytrees..making the model easier to write\n",
        "\n",
        "If we were using a MLP instead of a linear regressor, we might have millions of parameters. We're not going to create millions of arguments. For this reason, jax provides a very nice concept of [`pytrees`](https://jax.readthedocs.io/en/latest/pytrees.html). which is a container structure for the parameters. Specifically we will use a dictionary of params.\n",
        "\n",
        "Let's rewrite the model and run it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tdjqBuWnBUaU"
      },
      "source": [
        "def y(params, x):\n",
        "  return params['m'] * x + params['b']\n",
        "\n",
        "def mean_sqrd_loss(params, x, y_true):\n",
        "  y_pred = y(params, x)\n",
        "  loss = (y_pred - y_true) ** 2\n",
        "  return jnp.mean(loss)"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cCoWRC0n3JYL"
      },
      "source": [
        "Now function transforms like `grad` are `pytree` aware. They will return gradients with the same structure as the input params"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pgRgcL4CwTWq",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0f5a4747-b785-4caa-e5c0-e66ffd932bd6"
      },
      "source": [
        "true_params = {'m': m_true, 'b': b_true}\n",
        "model_params = {'m': 1.0, 'b': 0.0}\n",
        "\n",
        "# recall signature: mean_sqrd_loss(params, x, y_true)\n",
        "# note: grad takes gradients with respect to first arg\n",
        "grads = jit(grad(mean_sqrd_loss))\n",
        "\n",
        "print(\"grads w.r.t true_params\", grads(true_params, x, y_true))\n",
        "print(\"grads w.r.t model_params\", grads(model_params, x, y_true))"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "grads w.r.t true_params {'b': DeviceArray(-0.06965329, dtype=float32, weak_type=True), 'm': DeviceArray(0.8275, dtype=float32, weak_type=True)}\n",
            "grads w.r.t model_params {'b': DeviceArray(-4.069654, dtype=float32, weak_type=True), 'm': DeviceArray(170.11884, dtype=float32, weak_type=True)}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "auRIYnYXZW3I"
      },
      "source": [
        "## 3: parallelizing using pmap\n",
        "\n",
        "We can calculate gradients in parallel over single samples or over batches. These will need to be added for the gradient descent step.\n",
        " \n",
        "We'll use `pmap`. we will also reshape the input `(x, y)` data so we can `pmap` it as well. We won't `pmap` across the model params `m` or `b`. Each device will thur run the regression model, but with a fraction of the data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U75uXTIAYxaS",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "26c956da-7b88-499d-f46d-8ae9f83d4566"
      },
      "source": [
        "# reshape x and y_true to 8 sets of 16 examples\n",
        "shardsize = datasize//num_replicas\n",
        "x = x.reshape((num_replicas, shardsize))\n",
        "y_true = y_true.reshape((num_replicas, shardsize))\n",
        "x.shape, y_true.shape"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 16), (8, 16))"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XbvGd_gn7vhp"
      },
      "source": [
        "So we are going to have `num_replicas` sets of gradients, each of which corresponds to the `shardsize` samples that went to that specific device. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GcgbSkCv7n6o"
      },
      "source": [
        "from jax import pmap\n",
        "\n",
        "# recall signature: mean_sqrd_loss(params, x, y_true)\n",
        "p_grads = pmap(grad(mean_sqrd_loss), in_axes=(None, 0, 0))\n",
        "g = p_grads(model_params, x, y_true)"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qtuu-X6UIq0P"
      },
      "source": [
        "Notice the `in_axes` argument. Our loss signature is `mean_sqrd_loss(params, x, y_true)`. We dont wish to shard over params, but will shard using the devices dimension of both x and y.\n",
        "\n",
        "At this point we have the gradients for each of the `num_replicas` batches, and we will sum them and use the sum for an update."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gMs1WBAkI5K_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "101d618a-4d5e-4799-a699-f30b9097f41e"
      },
      "source": [
        "g = {k: jnp.sum(v) for k, v in g.items()}\n",
        "g"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'b': DeviceArray(-32.557228, dtype=float32),\n",
              " 'm': DeviceArray(1360.9507, dtype=float32)}"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j5g5FfJXJNYJ"
      },
      "source": [
        "jax provides `tree_map` to make this sum simpler and more versatile (all sorts of nesting is supported) "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hgRAZNR5JwPc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e743baa3-c656-41ed-c508-aa7e76a6af77"
      },
      "source": [
        "from jax.tree_util import tree_map \n",
        "\n",
        "g = p_grads(model_params, x, y_true)\n",
        "g = tree_map(jnp.sum, g)\n",
        "g"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'b': DeviceArray(-32.557228, dtype=float32),\n",
              " 'm': DeviceArray(1360.9507, dtype=float32)}"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "faOX8MgbJM6O"
      },
      "source": [
        "There are a number of parallel ops we can call that operate _across_ the devices. These include soncepts such as `allreduce`. After all we want to combine sharded data in different patterns. Here we use `psum`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yYbFKCcPFtjd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e9353122-db0f-42d8-e222-bfae5d565688"
      },
      "source": [
        "from jax.lax import psum\n",
        "\n",
        "def p_grads(params, x, y_true):\n",
        "  grads = jit(grad(mean_sqrd_loss))(params, x, y_true)\n",
        "  grads = tree_map(lambda v: psum(v, 'device'), grads)\n",
        "  return grads\n",
        "\n",
        "p_grads = pmap(p_grads, in_axes=(None, 0, 0), axis_name='device')\n",
        "\n",
        "model_params = {'m': 1.0, 'b': 0.0}\n",
        "p_grads(model_params, x, y_true)"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'b': ShardedDeviceArray([-32.557228, -32.557228, -32.557228, -32.557228,\n",
              "                     -32.557228, -32.557228, -32.557228, -32.557228],                   dtype=float32),\n",
              " 'm': ShardedDeviceArray([1360.9507, 1360.9507, 1360.9507, 1360.9507, 1360.9507,\n",
              "                     1360.9507, 1360.9507, 1360.9507], dtype=float32)}"
            ]
          },
          "metadata": {},
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice how the `psum` did an allreduce on each machine by dint of being within the `pmap`, so that the gradient was replicated on the local shards on the devices. Also note how we baned the \"devine\" axis and used it to make sure that the allredice was done over the same axes.\n",
        "\n",
        "Now its time to uodate the model step by step. And this can be donr using `tree_multimap` within the pmap call too. Since we allreduce the fradients on each device, we'll do the updates on each device as well."
      ],
      "metadata": {
        "id": "orJ_QHpDYCUQ"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7v0N3huqMWxd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f892bcdd-0158-4fc0-9e2d-5abd757a3945"
      },
      "source": [
        "from jax.tree_util import tree_multimap\n",
        "\n",
        "def p_update(params, x, y_true):\n",
        "  # calculate gradients summing across all devices\n",
        "  grads = jit(grad(mean_sqrd_loss))(params, x, y_true)\n",
        "  grads = tree_map(lambda v: psum(v, 'device'), grads)  \n",
        "  # apply update\n",
        "  def update(p, g):  \n",
        "    learning_rate = 0.001\n",
        "    return p - learning_rate * g\n",
        "  new_params = tree_multimap(update, params, grads)\n",
        "  # return new params\n",
        "  return new_params\n",
        "\n",
        "p_update = pmap(p_update, in_axes=(None, 0, 0), axis_name='device')\n",
        "\n",
        "model_params = {'m': 1.0, 'b': 0.0}\n",
        "p_update(model_params, x, y_true)"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:189: FutureWarning: jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() instead as a drop-in replacement.\n",
            "  'instead as a drop-in replacement.', FutureWarning)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'b': ShardedDeviceArray([0.03255723, 0.03255723, 0.03255723, 0.03255723,\n",
              "                     0.03255723, 0.03255723, 0.03255723, 0.03255723],                   dtype=float32),\n",
              " 'm': ShardedDeviceArray([-0.3609507, -0.3609507, -0.3609507, -0.3609507,\n",
              "                     -0.3609507, -0.3609507, -0.3609507, -0.3609507],                   dtype=float32)}"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeBq8a6mNzer"
      },
      "source": [
        "Now notice that the updated params are sharded as well. So we might as well have provided the initial params aharded, and get rid of the `None` for the params dimension. This will enable us to put this routine into a nice for loop."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Exk_R0uBM0MU",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ed1ce886-8ed1-45cc-fed4-9b418e82364c"
      },
      "source": [
        "def shard(x):\n",
        "  return pmap(lambda v: v)(x)\n",
        "\n",
        "def replicate(x, replicas=num_replicas):  \n",
        "  replicated = jnp.stack([x] * replicas)\n",
        "  return shard(replicated)\n",
        "\n",
        "model_params = {'m': 1.0, 'b': 0.0}\n",
        "model_params = tree_map(replicate, model_params)\n",
        "model_params"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'b': ShardedDeviceArray([0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32, weak_type=True),\n",
              " 'm': ShardedDeviceArray([1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32, weak_type=True)}"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sZIYAajm4Yv6"
      },
      "source": [
        "If we explicitly shard x and y too, instead of reshaping them, we'll also start by shipping parts of the data across the devices"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cVbER2R74XGq",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5cc996ba-57d3-46d8-d528-335891b47d03"
      },
      "source": [
        "x = shard(x)\n",
        "y_true = shard(y_true)\n",
        "  \n",
        "x.shape, type(x)"
      ],
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 16), jax.interpreters.pxla._ShardedDeviceArray)"
            ]
          },
          "metadata": {},
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EF9KDsyH4tPa"
      },
      "source": [
        "Now we can run smoothly in a loop!\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ve68pTe81hgC",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d48bf3c0-436d-45a9-92a9-32f6eac3b551"
      },
      "source": [
        "from jax import value_and_grad\n",
        "\n",
        "def p_update(params, x, y_true):\n",
        "  # calculate gradients\n",
        "  loss, grads = jit(value_and_grad(mean_sqrd_loss))(params, x, y_true)\n",
        "  # do all reduce sum across all replicas\n",
        "  grads = tree_map(lambda v: psum(v, 'device'), grads)  \n",
        "  # apply update\n",
        "  def update(p, g):  \n",
        "    learning_rate = 0.001\n",
        "    return p - learning_rate * g\n",
        "  new_params = tree_multimap(update, params, grads)\n",
        "  return new_params, loss\n",
        "\n",
        "# note: we pmap across _all_ args now; model_params, x and y\n",
        "p_update = pmap(p_update, in_axes=(0, 0, 0), axis_name='device')\n",
        "\n",
        "for i in range(20):\n",
        "  model_params, loss = p_update(model_params, x, y_true)  \n",
        "  print(\"step\", i , \"loss\", jnp.mean(loss), \n",
        "        \"m\", model_params['m'][0] , \"b\", model_params['b'][0])"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/jax/_src/tree_util.py:189: FutureWarning: jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() instead as a drop-in replacement.\n",
            "  'instead as a drop-in replacement.', FutureWarning)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "step 0 loss 221.33818 m -0.3609507 b 0.03255723\n",
            "step 1 loss 52.395218 m -0.98463047 b 0.06459354\n",
            "step 2 loss 16.815964 m -1.2704427 b 0.09611727\n",
            "step 3 loss 9.247494 m -1.4014213 b 0.12713662\n",
            "step 4 loss 7.564609 m -1.4614446 b 0.15765966\n",
            "step 5 loss 7.1207185 m -1.4889513 b 0.18769434\n",
            "step 6 loss 6.9399004 m -1.5015568 b 0.21724845\n",
            "step 7 loss 6.817111 m -1.5073334 b 0.24632971\n",
            "step 8 loss 6.7092013 m -1.5099807 b 0.27494568\n",
            "step 9 loss 6.607023 m -1.5111939 b 0.30310377\n",
            "step 10 loss 6.5085716 m -1.5117497 b 0.33081135\n",
            "step 11 loss 6.4133463 m -1.5120045 b 0.3580756\n",
            "step 12 loss 6.3211675 m -1.5121213 b 0.3849036\n",
            "step 13 loss 6.2319174 m -1.5121748 b 0.4113024\n",
            "step 14 loss 6.1455016 m -1.5121993 b 0.43727878\n",
            "step 15 loss 6.0618305 m -1.5122105 b 0.46283954\n",
            "step 16 loss 5.9808154 m -1.5122156 b 0.48799133\n",
            "step 17 loss 5.9023705 m -1.512218 b 0.5127407\n",
            "step 18 loss 5.826417 m -1.5122191 b 0.53709406\n",
            "step 19 loss 5.752874 m -1.5122195 b 0.5610578\n"
          ]
        }
      ]
    }
  ]
}