dfm/intro-to-jax-part2.ipynb

## intro-to-jax-part2.ipynb
{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/dfm/6be59481367fe4c48081089e97e7047e/intro-to-jax-part2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ada14b6a-d989-4aa4-8b71-0d870933eb13",
      "metadata": {
        "id": "ada14b6a-d989-4aa4-8b71-0d870933eb13"
      },
      "source": [
        "# Introduction to JAX (Part 2)\n",
        "\n",
        "We'll start with a re-cap of the previous \"intro to jax\" session with (hopefully!) enough info to get people who weren't there caught up.\n",
        "\n",
        "This tutorial includes a whirlwind introduction to JAX. It's going to be pretty incomplete so, if you want more info, check out the [JAX docs](https://jax.readthedocs.io).\n",
        "\n",
        "We'll pretty much always want to include this line since JAX normally operates with single point precision:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "9dc1746e-58c1-4ae6-824c-f0e0865a784a",
      "metadata": {
        "id": "9dc1746e-58c1-4ae6-824c-f0e0865a784a"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "\n",
        "# In many cases you may want to enable support for double precision\n",
        "# jax.config.update(\"jax_enable_x64\", True)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b2f1843d-fe6d-428d-9206-08599ce547de",
      "metadata": {
        "id": "b2f1843d-fe6d-428d-9206-08599ce547de"
      },
      "source": [
        "## `jax.numpy`\n",
        "\n",
        "`jax.numpy` works just like `numpy` (almost always):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "467fe662-819a-4c5a-b179-8e5d1cb21076",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "467fe662-819a-4c5a-b179-8e5d1cb21076",
        "outputId": "2c7c7145-e1f6-445b-eea4-d4cc158509cf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray([0.1  , 1.325, 2.55 , 3.775, 5.   ], dtype=float32),\n",
              " DeviceArray([ 0.09983342,  0.9699439 ,  0.55768377, -0.5918946 ,\n",
              "              -0.9589243 ], dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "import jax.numpy as jnp\n",
        "\n",
        "x = jnp.linspace(0.1, 5.0, 5)\n",
        "y = jnp.sin(x)\n",
        "x, y"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e0b320f6-d4fd-4f59-bf0a-8cbd748215be",
      "metadata": {
        "id": "e0b320f6-d4fd-4f59-bf0a-8cbd748215be"
      },
      "source": [
        "We can combine regular `numpy` and `jax.numpy`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "8e108d6f-5bf7-483f-8485-7bab4533f92f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8e108d6f-5bf7-483f-8485-7bab4533f92f",
        "outputId": "cdfd2bd5-227a-4dcc-be5d-19139a0aa104"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(array([0.1  , 1.325, 2.55 , 3.775, 5.   ]),\n",
              " DeviceArray([ 0.09983342,  0.9699439 ,  0.55768377, -0.5918946 ,\n",
              "              -0.9589243 ], dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "x = np.linspace(0.1, 5.0, 5)\n",
        "y = jnp.sin(x)\n",
        "x, y"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9282cf37-e89e-4f3a-abfc-b247538f6e4b",
      "metadata": {
        "id": "9282cf37-e89e-4f3a-abfc-b247538f6e4b"
      },
      "source": [
        "## `jax.jit`\n",
        "\n",
        "We use `jax.jit` to fuse operations, and run them on the GPU, for example.\n",
        "One of the key points to remember when using JAX is that it works best in a \"functional\" style.\n",
        "A lot of the key JAX functions take a function as input and return a new function.\n",
        "For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "cc4dc0ac-b240-47d1-91d3-f1edbb0a0526",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cc4dc0ac-b240-47d1-91d3-f1edbb0a0526",
        "outputId": "f078464d-67f4-43d8-8983-8d2c2acf5e91"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hi from this function\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([2.6049867, 4.1377964, 3.246622 , 2.053278 , 1.883305 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "def jnp_function(x):\n",
        "    print(\"hi from this function\")\n",
        "    arg = jnp.sin(x)\n",
        "    return 1.5 + jnp.exp(arg)\n",
        "\n",
        "jitted_function = jax.jit(jnp_function)\n",
        "\n",
        "jitted_function(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8c7d5f64-7797-46d8-ac54-4c46cb6d9525",
      "metadata": {
        "id": "8c7d5f64-7797-46d8-ac54-4c46cb6d9525"
      },
      "source": [
        "What happens if we call that function again?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "79bac88b-202b-42b2-889c-16afe755f667",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "79bac88b-202b-42b2-889c-16afe755f667",
        "outputId": "4c21b7ef-c766-4415-ce9f-78757d165879"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([2.6049867, 4.1377964, 3.246622 , 2.053278 , 1.883305 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "jitted_function(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e7c75a4-23da-4c28-9ee0-1e5baa960327",
      "metadata": {
        "id": "5e7c75a4-23da-4c28-9ee0-1e5baa960327"
      },
      "source": [
        "What about if we call it with a different input?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "35573796-3a57-44ab-94f6-dcd34c481f97",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "35573796-3a57-44ab-94f6-dcd34c481f97",
        "outputId": "6b5c0dea-ea02-493e-ac8c-1d7592af2fb9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([2.6048036, 3.7815475, 3.1976116, 2.0723903, 1.9410601], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "jitted_function(np.sin(x))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "020331f8-2033-4937-8273-7565ab26fd6d",
      "metadata": {
        "id": "020331f8-2033-4937-8273-7565ab26fd6d"
      },
      "source": [
        "What about an input with a different shape?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "1353bc5c-6a1d-4d12-a9be-28e9204e2c95",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1353bc5c-6a1d-4d12-a9be-28e9204e2c95",
        "outputId": "7be7f76c-4f73-405c-8a0b-945d23ba105f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hi from this function\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([2.6049867, 4.1377964, 3.246622 , 2.053278 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ],
      "source": [
        "jitted_function(x[:-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b2990af6-a1b8-426a-aa5e-7560c97b64d6",
      "metadata": {
        "id": "b2990af6-a1b8-426a-aa5e-7560c97b64d6"
      },
      "source": [
        "*Note:* It is common to use `jax.jit` as a \"decorator\":"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "d997f1d0-8117-425b-bd8a-889be7cb821f",
      "metadata": {
        "id": "d997f1d0-8117-425b-bd8a-889be7cb821f"
      },
      "outputs": [],
      "source": [
        "@jax.jit\n",
        "def jitted_function(x):\n",
        "    arg = jnp.sin(x)\n",
        "    return 1.5 + jnp.exp(arg)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "What about control flow?"
      ],
      "metadata": {
        "id": "Se90GwOHOhP3"
      },
      "id": "Se90GwOHOhP3"
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def incorrect_conditional_func(x):\n",
        "    if jnp.all(x > 0):\n",
        "        return x\n",
        "    arg = jnp.sin(x)\n",
        "    return 1.5 + jnp.exp(arg)"
      ],
      "metadata": {
        "id": "XG8W5MNUOjkk"
      },
      "id": "XG8W5MNUOjkk",
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# What happens if we run this?\n",
        "# incorrect_conditional_func(x)"
      ],
      "metadata": {
        "id": "p5ZfoQGvOopU"
      },
      "id": "p5ZfoQGvOopU",
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def correct_conditional_func(x):\n",
        "    arg = jnp.sin(x)\n",
        "    return jnp.where(jnp.all(x > 0), x, 1.5 + jnp.exp(arg))"
      ],
      "metadata": {
        "id": "iIIaFa_1O88z"
      },
      "id": "iIIaFa_1O88z",
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "correct_conditional_func(x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5HNdGm6gPMZ_",
        "outputId": "12092eb1-8b97-47a4-c827-0972934d2422"
      },
      "id": "5HNdGm6gPMZ_",
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([0.1  , 1.325, 2.55 , 3.775, 5.   ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cdc4a86a-cd09-4d4d-a5d0-55438cc9c02b",
      "metadata": {
        "id": "cdc4a86a-cd09-4d4d-a5d0-55438cc9c02b"
      },
      "source": [
        "## `jax.vmap`\n",
        "\n",
        "`jax.vmap` gives a mechanism for applying a \"scalar\" function on a vector of inputs.\n",
        "The same effects can often be achieved by manually broadcasting, but `vmap` comes in handy more often than you might think."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "136772d8-a384-4dc3-87ba-e967e5bfbf7c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "136772d8-a384-4dc3-87ba-e967e5bfbf7c",
        "outputId": "e5eb8a3d-e696-4880-94e9-5074c315e837"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[[ 1.19428426e-01,  2.83938229e-01,  1.14193819e-01],\n",
              "              [ 2.83938229e-01,  6.75056338e-01,  2.71493077e-01],\n",
              "              [ 1.14193819e-01,  2.71493077e-01,  1.09188654e-01]],\n",
              "\n",
              "             [[ 1.69821870e+00, -1.17982101e+00, -5.81696212e-01],\n",
              "              [-1.17982101e+00,  8.19669247e-01,  4.04127836e-01],\n",
              "              [-5.81696212e-01,  4.04127836e-01,  1.99250251e-01]],\n",
              "\n",
              "             [[ 2.88318753e-01, -3.12033236e-01, -1.95758328e-01],\n",
              "              [-3.12033236e-01,  3.37698251e-01,  2.11859629e-01],\n",
              "              [-1.95758328e-01,  2.11859629e-01,  1.32913038e-01]],\n",
              "\n",
              "             [[ 8.65139291e-02,  8.35990533e-03,  1.60806060e-01],\n",
              "              [ 8.35990533e-03,  8.07823846e-04,  1.55388089e-02],\n",
              "              [ 1.60806060e-01,  1.55388089e-02,  2.98895091e-01]],\n",
              "\n",
              "             [[ 5.42364597e-01,  1.19975701e-01,  3.55058730e-01],\n",
              "              [ 1.19975701e-01,  2.65396535e-02,  7.85420388e-02],\n",
              "              [ 3.55058730e-01,  7.85420388e-02,  2.32439041e-01]]],            dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "A = np.random.default_rng(1).normal(size=(5, 3))\n",
        "\n",
        "def scalar_function(x):\n",
        "    return jnp.outer(x, x)\n",
        "\n",
        "vector_function = jax.vmap(scalar_function)\n",
        "vector_function(A)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "4a861a6d-f6b2-415b-a3d7-a9eba66b6df8",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4a861a6d-f6b2-415b-a3d7-a9eba66b6df8",
        "outputId": "cb2d525b-6174-4e9d-818c-61077554e40b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[[ 1.19428434e-01,  2.83938242e-01,  1.14193830e-01],\n",
              "        [ 2.83938242e-01,  6.75056374e-01,  2.71493097e-01],\n",
              "        [ 1.14193830e-01,  2.71493097e-01,  1.09188661e-01]],\n",
              "\n",
              "       [[ 1.69821877e+00, -1.17982104e+00, -5.81696252e-01],\n",
              "        [-1.17982104e+00,  8.19669245e-01,  4.04127838e-01],\n",
              "        [-5.81696252e-01,  4.04127838e-01,  1.99250259e-01]],\n",
              "\n",
              "       [[ 2.88318777e-01, -3.12033246e-01, -1.95758328e-01],\n",
              "        [-3.12033246e-01,  3.37698251e-01,  2.11859620e-01],\n",
              "        [-1.95758328e-01,  2.11859620e-01,  1.32913032e-01]],\n",
              "\n",
              "       [[ 8.65139256e-02,  8.35990480e-03,  1.60806056e-01],\n",
              "        [ 8.35990480e-03,  8.07823801e-04,  1.55388084e-02],\n",
              "        [ 1.60806056e-01,  1.55388084e-02,  2.98895090e-01]],\n",
              "\n",
              "       [[ 5.42364622e-01,  1.19975697e-01,  3.55058738e-01],\n",
              "        [ 1.19975697e-01,  2.65396512e-02,  7.85420322e-02],\n",
              "        [ 3.55058738e-01,  7.85420322e-02,  2.32439032e-01]]])"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "A[:, None, :] * A[:, :, None]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a37c40b8-8a99-46db-860b-a04b2918b976",
      "metadata": {
        "id": "a37c40b8-8a99-46db-860b-a04b2918b976"
      },
      "source": [
        "## `jax.grad`\n",
        "\n",
        "Any JAX function can also be differentiated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "fc0e710d-962b-44e7-8649-51b88e47c806",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fc0e710d-962b-44e7-8649-51b88e47c806",
        "outputId": "501f0dd5-9aee-4df1-c3d1-ff14e7e8bf15"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(1.4174242, dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ],
      "source": [
        "grad_function = jax.grad(jitted_function)\n",
        "grad_function(0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fc51106f-1379-4b53-95f1-031a5d67264d",
      "metadata": {
        "id": "fc51106f-1379-4b53-95f1-031a5d67264d"
      },
      "source": [
        "By default, differentiation is only supported for scalar outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "c343dc80-4d7e-461b-b520-34dfadff76f2",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c343dc80-4d7e-461b-b520-34dfadff76f2",
        "outputId": "eb595fe6-4fda-4bc1-a331-b81f18737388"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "    with pytest.raises(TypeError) as info:\n",
            ">       jax.grad(jitted_function)(x)\n",
            "E       TypeError: Gradient only defined for scalar-output functions. Output had shape: (5,).\n",
            "\n",
            "<ipython-input-16-b967da8725c7>:4: TypeError\n"
          ]
        }
      ],
      "source": [
        "import pytest\n",
        "\n",
        "with pytest.raises(TypeError) as info:\n",
        "    jax.grad(jitted_function)(x)\n",
        "print(\"\\n\\n\".join(str(info.getrepr()).split(\"\\n\\n\")[-2:]))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6d2a14ce-d5bd-436f-801a-71242d1167f9",
      "metadata": {
        "id": "6d2a14ce-d5bd-436f-801a-71242d1167f9"
      },
      "source": [
        "But we can combine `grad` with `vmap` to get the derivative at each input point:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "59edef16-1a86-4ce4-8ac6-5e470090c242",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "59edef16-1a86-4ce4-8ac6-5e470090c242",
        "outputId": "fc867951-c5f6-4e37-88d2-3b98353b8767"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 1.0994664 ,  0.6418517 , -1.4497899 , -0.4459506 ,\n",
              "              0.10872913], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "jax.vmap(jax.grad(jitted_function))(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fb4d5bbf-4e23-4cd9-be05-d21bf89094fe",
      "metadata": {
        "id": "fb4d5bbf-4e23-4cd9-be05-d21bf89094fe"
      },
      "source": [
        "Another useful function is `jax.value_and_grad`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "id": "ce76be99-0778-4859-90cf-6fce4c1cf3cd",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ce76be99-0778-4859-90cf-6fce4c1cf3cd",
        "outputId": "dc0ae4a5-810d-42f3-da39-e5a4a0623bec"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray([2.6049867, 4.1377964, 3.246622 , 2.053278 , 1.883305 ], dtype=float32),\n",
              " DeviceArray([ 1.0994664 ,  0.6418517 , -1.4497899 , -0.4459506 ,\n",
              "               0.10872913], dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ],
      "source": [
        "jax.vmap(jax.value_and_grad(jitted_function))(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "id": "ba508356-266d-49cb-b6a9-9fbc46c0b7bf",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "ba508356-266d-49cb-b6a9-9fbc46c0b7bf",
        "outputId": "7c7d4b42-2c95-4585-d06f-5df70b7d3db1"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXIAAAD4CAYAAADxeG0DAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3dd3iUVdrH8e9JI6SQhCRAIJVOEmpCR0UpinQQEMGCIujay9pXsK1r17W8iKCCgIgUQUEQVhQFEQg1oVcJNaGEJJA65/3jYKdnZp4p9+e6ckFgMs9vYHLnPKcqrTVCCCHcl4/VAYQQQlSMFHIhhHBzUsiFEMLNSSEXQgg3J4VcCCHcnJ8VF42KitKJiYlWXFoIIdxWRkZGrtY6+q9/bkkhT0xMZNWqVVZcWggh3JZSas+Z/ly6VoQQws1JIRdCCDcnhVwIIdycJX3kQghxPqWlpWRnZ1NUVGR1FKcLDAwkNjYWf3//C3q8FHIhhEvKzs4mNDSUxMRElFJWx3EarTVHjhwhOzubpKSkC/oa6VoRQrikoqIiIiMjvaqIAyiliIyMvKg7ESnkQgiX5W1F/FcX+7o9omul3KbZcjCfA3mnOJBXREFxGQ2qh9I4NoyokEpWxxPCoXLyi9l6KJ8tB/OxaU1CZDCJkUEkRAYT4CdtNW/g1oX8SEExU1fuZcrPv7Dv+KkzPiauamVuaZfEkNbxBPr7OjmhEI5RWm5j1pp9vP/9DnbkFJ7xMRFB/gxtk8CNbROoFhro5ITeKSQkhIKCAqdf1y0LeUmZjTcXbWXcD7soKbfRrk4kD19dn6SoEGLCAgn092XzgRNs2JfHok2HeO6rjYxdsoO7rqzLDa3i8fOVVopwT1prPlu5l7e/3c6+46dIjqnCU90b0SimCvWrh+Lno9hz9CS7cwuZu+EA7yzezvvf72RQyzge69aQ4Epu+S0vzsPt/ld35BRw/9S1bNiXR7/mtbizYx3qVQ/92+Na146kde1Ihl9Wm592HOGNhVt5enYW8zMP8vbg5kRKl4twM/lFpTw6Yz3zNhykRXw4z/dJpWOD6L/1p0YEB9AsLpw+zWuxK7eQcT/sZNLPe1i6I5d3b2hBo5gqFr0C9/PYY48RFxfHXXfdBcDo0aPx8/Nj8eLFHDt2jNLSUp5//nl69+79p6/77rvvePXVV/nqq68AuPvuu0lPT+eWW24hIyODBx98kIKCAqKiovj444+JiYmpUE63KuTTVu5l1JwsKvn7MGZoC65JvbAX37ZOJG1qt2F6RjZPfpFJr3eWMmZoGo1jwxycWAj72HzwBP+YtJo9R0/yeLeGjLi89gUNiCVFBfNC38Z0bxLD/VPX0vvdpYzqmcyQ1glOSG0/z3yZxcb9J+z6nMk1qzCqZ8o5HzNo0CDuv//+3wr5tGnTWLBgAffeey9VqlQhNzeXNm3a0KtXrwv6/ygtLeWee+5h9uzZREdH89lnn/Hkk0/y4YcfVui1uFUhzy8uIy0hglcHNKVG2MX1+SmlGJAeR8MaVbhjUgb9xyzjvRta0Dm5uoPSCmEfGXuOctP4FQRX8mPK8Na0rh150c/Rrk4U8+67jIemrePJWZkcP1nKXVfWdUBaz9K8eXMOHz7M/v37ycnJISIigho1avDAAw+wZMkSfHx82LdvH4cOHaJGjRrnfb4tW7aQmZlJly5dACgvL69waxzsWMiVUr7AKmCf1rqHvZ73j4a1S2RYu0R8fC59SlLj2DC+vKcDwz5eyT8mr2b8LelcVu9vu0IK4RLW7j3OzR+upHqVQD4d0YbqVS590DIqpBLjb07noc/X8cqCLQBuU8zP13J2pAEDBjB9+nQOHjzIoEGDmDx5Mjk5OWRkZODv709iYuLf5nz7+flhs9l++/zXv9dak5KSwk8//WTXjPYc9bsP2GTH5/sbHx9VoSL+q6rBAUwc1oo61UK4feIqVuw6aod0QtjXhuw8bhz/M5EhAUy5vWJF/Fd+vj68NqApvZvV5JUFW3h38XY7JPVsgwYNYurUqUyfPp0BAwaQl5dHtWrV8Pf3Z/HixezZ8/edZRMSEti4cSPFxcUcP36c//3vfwA0aNCAnJyc3wp5aWkpWVlZFc5ol0KulIoFugPj7PF8zhAW5M8nt7WiZnhlbv14JZn78qyOJMRv9hwp5MYPfyassj9Tbm9z0V2J5/LXYj5zdbbdntsTpaSkkJ+fT61atYiJiWHIkCGsWrWKxo0bM3HiRBo2bPi3r4mLi2PgwIGkpqYycOBAmjdvDkBAQADTp0/n0UcfpWnTpjRr1oxly5ZVOKPSWlf8SZSaDrwIhAIPn6lrRSk1AhgBEB8fn3amn2JWOJhXRP//W4bWmi/v6SCzWYTlCovL6PfeMg7lFzH7rvYkRAY75Dql5TZuHP8zq385zrSRbWkWF+6Q61yqTZs20ahRI6tjWOZMr18plaG1Tv/rYyvcIldK9QAOa60zzvU4rfVYrXW61jo9Otp1+qRrhAUyZmgauYUl3PPpGsrKbef/IiEcRGvNQ9PWse1wPu8MbuGwIg7g7+vDe0PSqBZaiRETV3HohPftMugp7NG10h7opZTaDUwFrlJKTbLD8zpN49gw/t23Mct2HOHl04NAQljh3cXbmZ91kCeubUSHelEOv17V4ADG3ZxOQXEZIyauoris3OHXFPZX4UKutX5cax2rtU4Erge+1VoPrXAyJ7suLZab2iYwdslO5m04YHUc4YWWbc/ltYVb6dOsJrd1uLDtS+2hYY0qvD6wKeuy83hlvjRk3JGsVf+Dp7on0ywunMdmrOdA3pn3bhHCEfJOlvLgtHUkRQbz736Nnb7r3zWpMQxtE8+4H3exZGuOU68tKs6uhVxr/Z2j5pA7Q4CfD28OakZpuebhz9dhs1V8IFiI89Fa88QXG8gtKOat65sTFGDNOr0nr02mbrUQHvp8HUcLSyzJIC6NtMj/IjEqmH/1SGbp9iN8vGy31XGEF5i5eh9z1x/ggS71Ld02onKAL/+9vjl5J0t5ZPp67DGjTTiHFPIzGNwqjk4Nq/Gf+ZvZeijf6jjCg2UfO8moOVm0SqzKHVfUsToOyTWr8Mg1DVi06RDTM2R+ub0lJiaSm5tr9+eVQn4GSileuq4JoZX8ePjzdZRLF4twAK01T8zKxKY1rw1siq8dVi3bw63tk2iZGMFzX23ksExJPK+ysjKrI0ghP5uokEqM7pXC+uw8Plq6y+o4wgPNWrOPJVtzeOTqBsRVDbI6zm98fBT/6d+EojIbT8+u+PJxd/fcc8/RoEEDOnTowODBg3n11Vfp2LEj999/P+np6bz11lt8+eWXtG7dmubNm9O5c2cOHToEwJEjR+jatSspKSkMHz7cYd1VbrX7obP1aBLD7LX7ePWbLXRNrkF8pOt8swn3lltQzLNfbSQtIYIb2yZaHedv6kSH8EDn+rw0fzPzNhzg2sYV36GvQr5+DA5usO9z1mgM3f5zzoesXLmSGTNmsG7dOkpLS2nRogVpaWkAlJSUsGrVKgCOHTvG8uXLUUoxbtw4Xn75ZV577TWeeeYZOnTowNNPP83cuXMZP368fV/DadIiPwelFM/1ScXPx4cnZm2QwR9hN6PnZHGyuJyX+jd2mS6Vv7r9siQa1wrj6dmZHPPSWSxLly6ld+/eBAYGEhoaSs+ePX/7u0GDBv32++zsbK6++moaN27MK6+88ttGWEuWLGHoULOspnv37kRERDgkp7TIzyMmrDKPdmvIv77IZHpGNgPS46yOJNzc/zYd4qv1B3ioS33qVvv76Vauws/Xh5f6N6HnOz/y0vzN/Kd/E+vCnKflbIXg4N+3T7jnnnt48MEH6dWrF9999x2jR492ahZpkV+AIa3iSU+I4N/zNnlty0TYx6mScp6enUW9aiGMdIFZKueTXLMKwzskMXXlXlbt9r7tntu3b8+XX35JUVERBQUFvx3d9ld5eXnUqlULgAkTJvz255dffjlTpkwB4Ouvv+bYsWMOySmF/AL4+Cie75vKiaIy/vP1ZqvjCDf232+3se/4KZ7vk0qAn3t8+93XuR61wivz5KxMSr1sU7mWLVvSq1cvmjRpQrdu3WjcuDFhYX+f6z969GgGDBhAWloaUVG/75EzatQolixZQkpKCjNnziQ+Pt4xQbXWTv9IS0vT7uiFuRt1wqNf6VW7j1gdRbihLQdP6DqPz9UPTVtrdZSL9k3WQZ3w6Fd6zHfbnXbNjRs3Ou1a55Kfn6+11rqwsFCnpaXpjIwMp1z3TK8fWKXPUFPdo0ngIu7rVI+aYYE8OStTtrsVF0VrzVNfZBIS6Mfj3f5+EIGr65JcnS7J1Xlz0Tayj520Oo5TjRgxgmbNmtGiRQv69+9PixYtrI70N1LIL0JwJT+e7pnC5oP5snxfXJSZq/exYtdRHrumodseXjK6lzk387mvNlqcxLmmTJnC2rVr2bx5M48//rjVcc5ICvlFujqlOlc2iOaNhVtlI35xQfJOlfLi15toHh/OQDee9VQrvDL3dKrLgqxDLN5y2CnX1F465fdiX7cU8ouklGJ0rxRKbZoX5jr0rGnhId5YuJWjhSU81zvVLoeHW2l4h9rUjg5m9JwsikodewhFYGAgR44c8bpirrXmyJEjBAZe+DmtMo/8EiREBnPnFXV463/buL5VHO3qOP4kF+GesvbnMfGn3Qxtk0BqLet2NrSXAD8fnumVwo3jVzB2yU7u7VTPYdeKjY0lOzubnBzv2x89MDCQ2NjYC368FPJLdGfHOsxck82o2VnMu+8y/H3l5kb8mc2mGTU7i4igAB7q0sDqOHZzWb1oujeO4d3F2+nbvJbD9onx9/cnKcl5JyW5M6k+lyjQ35dRPVLYdrhANtUSZzRzzT5W7TnGo90aEhbkb3Ucu3qqRyN8fRTPfOldA5+uSgp5BXROrk6nhtV4a9E2DubJwKf4Xd7JUl6ct4kW8eFc1+LCb5HdRUxYZe7tVI9Fmw7x7eZDVsfxelLIK2hUTzPw+fxcaZmI3722cAvHTpbwrAcMcJ7Nre2TqBMdzOg5Gx0+8CnOTQp5BcVHBvGPjnX4av0Blm63/8kfwv1k7stj0vI93OghA5xnE+Dnw7O9U/nl6Ene/36n1XG8mhRyO7jjijrEVw3i6dmZlJTJik9vZrNp/jU7k6rBATzY1XMGOM+mfd0oejSJ4b3vtvPLEe9a8elKpJDbQaC/L8/0SmFHTiEfysCnV/s8Yy9rfjnO490aEVbZswY4z+ap7sn4+ihGf5nldXO+XYUUcju5smE1uiZX561FZnc74X2OFpbw4tebaZkYQb8WtayO4zQ1wgJ5sEt9vt18mAVZMvBpBSnkdjTq9F4Uo+ScQ6/04rxNFBSV8ULfxijlmQOcZ3NLu0Qa1gjlmS+zKCy2/jBibyOF3I5qhVfmgS5mStY3WQetjiOcaMWuo3yekc3wy2pTv7rrnvrjKH6+PrzQN5UDeUW8uWir1XG8jhRyOxvWPomGNUIZPUdaJt6ipMzGU19soFZ4Ze7tVNfqOJZJS6jK9S3j+HDpbjYdOGF1HK9S4UKulApUSq1QSq1TSmUppZ6xRzB35X+6ZbI/r4g3FkrLxBuM+3EnWw8V8EyvFIICvHvXi0evaUhYZX8en7mBcpsMfDqLPVrkxcBVWuumQDPgGqVUGzs8r9tKS6jK4FbxfLRsNxuy86yOIxxoV24hby3axtUp1emcXN3qOJaLCA7g6R7JrN17nE9+2m11HK9R4UJ++gSigtOf+p/+8PofxY91a0hkcACPzFjvdeccegutNY/PXP/bwhhh9G5Wk8vrR/Pygi0yg8tJ7NJHrpTyVUqtBQ4DC7XWP5/hMSOUUquUUqu8YVvKsMr+PNs7lU0HTjB2iax680SfrdzL8p1HeeLaRlSvcuF7R3s6pRQv9ElFa3hq1gaZW+4EdinkWutyrXUzIBZopZT6W/NEaz1Wa52utU6Pjo62x2Vd3jWpNeiWWoO3/reNHTkF5/8C4TYOnSjihXmbaFPbDPCJP4urGsRDXeuzeEsOc9bttzqOx7PrrBWt9XFgMXCNPZ/XnT3TO4VAPx8en7EBmwz+eIRfD1IuKbPxYr8mXjdn/EINa59E07hwRs/JIie/2Oo4Hs0es1ailVLhp39fGegCbK7o83qKaqGB/KtHMit2H+UjObDZI8xas4+FGw/xUNf6JEUFWx3HZfn6KF4b0ITCknKelC4Wh7JHizwGWKyUWg+sxPSRf2WH5/UY16XF0qlhNV6ev5nth6WLxZ3tP36KUXOyaJkYwW0dalsdx+XVrRbKw13r883GQ3yxdp/VcTyWPWatrNdaN9daN9Fap2qtn7VHME+ilOLF/o0JCvDlwWlrZRaLm9Ja8+iM9ZSVa14d0BRfD91n3N5u61CbtIQIRs3OkgNYHERWdjpJtdBAnu/TmPXZeby3eIfVccQlmPzzL/ywLZcnujciIVK6VC6Ur4/i1QFNKSm38ciM9TJW5ABSyJ2oe5MYejWtydvfbmPNL8esjiMuwrZD+Tw/dyOX1YtiaOt4q+O4naSoYJ68thFLtubIVs8OIIXcyZ7rnUr1KoHcO3UNJ4pKrY4jLkBRaTn3fLqG4AA/XhvQVGapXKKhbRLoklydl+ZvJnOfrHi2JynkThYW5M9/Bzdn//EiHp8hI/nu4IW5m9h8MJ/XBjalmiz8uWRKKV7u34TI4Erc8+ka2VTOjqSQWyAtIYKHutZn7oYDfLpir9VxxDksyDrIJ8v3cPtlSXRsUM3qOG4vIjiANwY1Y/eRQp6eLScK2YsUcovccXkdLqsXxTNfZpG1X24zXdGu3EIe/nwdjWuF8c+rG1odx2O0rRPJPVfVY8bqbKas+MXqOB5BCrlFfHwUrw9sRkRQACM/yeBYYYnVkcQfFBaXMfKTVfj5KN4b0oIAP/lWsaf7OtXjivrRjJ6TJQP/diDvTgtFh1ZizI1pHM4v5u5PV1Mm88tdgtaaR2asZ/vhAt4e3IK4qkFWR/I4vj6Kt65vRo2wQO6ctJrcAlnCXxFSyC3WLC6c5/uksnT7EV6aLzsbuIL3l+xk7voDPHJNQzrUi7I6jscKDwpgzNA0jp0s4a7Jqykpk4bMpZJC7gIGpsdxc9sEPvhhF9NWyuCnleZnHuCl+Zvp3jiGkZfLEnxHS6kZxsvXNeHnXUd5bMZ6Gfy8RN59LpULeapHMjtzC3l81gaqhwVyRX3v2OrXlWTsOcZ9U9fSLC6c1wbKfHFn6d2sFnuOnOT1hVuJjwzi/s71rY7kdqRF7iL8fX14b0gLGlQP5R+TMmTBhJPtyi1k+ISVxIQFMu6mdAL9fa2O5FXuuaou16XF8uaibUzPyLY6jtuRQu5CQgP9+WhYS8KDAhj28Up+OXLS6khe4dCJIm75aAVKKT4e1orIkEpWR/I6Sin+3bcx7epE8tiM9SzceMjqSG5FCrmLqV4lkI+HtaS03MbgD5az96gUc0fKyS/mhg+Wk5tfzPib00mU/cUtE+Dnw/s3ppFSswp3TV7Nkq2efySkvUghd0H1qocy6bbW5BeVcsO45XKArYMcLSxh6Lif2X+8iI+GtaJ5fITVkbxeaKA/E25tRZ1qIYz4ZBXLdx6xOpJbkELuolJrhTFpeGuOnyzlhg+Ws1+KuV3lFhQzdNzP7D5SyPib02mVVNXqSOK08KAAPrmtFbERQdz68UqWbs+1OpLLk0LuwprEhjPx1lYcLSjhuv9bxvbD+VZH8gjZx04yYMxP7MwtYOxN6bSrK3PFXU1USCWmDG9NXEQQwz5ayfzMg1ZHcmlSyF1c8/gIpo5sQ0m55roxP5GxR5YzV8TWQ/n0/79lHCkoZtJtrWWapwurViWQz0a2IaVWFf4xOUPWWJyDFHI3kFIzjJl3tiOssj9Dxi1nQZa0Ti7Fj9tyGTDmJ7SGaXe0JT1RulNcXXhQAJOHt6Z93SgembGeVxZslhOGzkAKuZuIjwxi+h3taFA9lJGfZPD6wq3yhr5AWmvG/bCTmz78mRpVAplxZzsa1qhidSxxgYIC/Bh/c0uubxnHu4t3MOKTDApkL/M/kULuRqJDK/HZyLZclxbLf/+3jRGfrJJThs7jZEkZD32+jufnbqJrcg1m/qOdbILlhgL8fHixX2NG90xm8ZbD9HtvqYwZ/YEUcjcT6O/LK9c14dneKXy3JYdr3/qBFbuOWh3LJWXuy6PH2z8ya80+Huhcn/eGtCC4kuxK4a6UUtzSPokJw1qRW1BCj7d/ZMrPv8j+LEghd0tKKW5qm8i0O9ri66MYNPYnXpq/WXaPO63cphnz/Q76vreUk8XlTL6tNfd1roePj+yd4gk61Iti/n2X0TKxKk/M2sAdkzK8fhtcZcVPs/T0dL1q1SqnX9cTFRaX8dxXG5m6ci/1q4fwXO9UWteOtDqWZdZnH+fJWZls2JfHtY1r8O++jQkPCrA6lnAAm00z7sedvLJgC0EBfjxxbUMGpMV59A9spVSG1jr9b38uhdwzfLv5EP/6Iot9x0/Rr3ktHr+2EdGh3rNnyPGTJbyxcCsTl+8hKqQST/dIpkeTGNnB0AtsP5zPEzMzWbH7KK2SqjKqZzIpNcOsjuUQDivkSqk4YCJQHdDAWK31W+f6GinkjnGqpJx3Fm9j7JKd+Pv6cFuHJIZfVpuwyv5WR3OYkyVlfLR0N2O+30FBcRk3tUngoasbUCXQc1+z+DubTTM9I5t/f72J4ydL6dOsJg91beBxA9uOLOQxQIzWerVSKhTIAPporTee7WukkDvWzpwCXlu4lbnrDxBW2Z/hHZIY2iaBiGDP6WLILyrls5V7GbtkJ4fzi+nUsBoPX92ARjEyrdCb5Z0q5f3vd/Dh0l2U2zT9W8Qy4vLa1I4OsTqaXTita0UpNRt4R2u98GyPkULuHJn78njtmy0s3pJDZX9frkuL5Zb2idRx4zf13qMnmbR8D1N+/oX84jLa1K7KQ10b0FIW94g/OJhXxDuLtzFtVTal5Ta6JlfnlnZJtKld1a2725xSyJVSicASIFVrfeIvfzcCGAEQHx+ftmfPHrtdV5zb5oMn+PDHXXyxZj8l5TbSEyK4Li2W7k1iCHWDLohTJeV8s/Egn63cy7IdR/BR0L1JTW6/LIkmseFWxxMuLCe/mAnLdjPxp92cKCojMTKIgS3j6Nu8FjFhla2Od9EcXsiVUiHA98ALWuuZ53qstMitkZNfzIzV2Xy+ai87cgoJ8POhQ90ouiZXp1Oj6i41OHqssITvth5mfuZBvt+aQ1GpjbiqlRmQFkf/tFhqhbvfN6GwzqmScr7OPMDUlXt/W3fRLC6cbqk16NSoOnWig92ipe7QQq6U8ge+AhZorV8/3+OlkFtLa82avceZu/4AC7IOkn3MbJHboHoobetE0qZ2VZrEhhMTFui0N/fhE0Wsy85j5e6jLN2ey8YDJ9AaqlepxNUpNeiWGkPrpKoePbVMOMfu3ELmZR7g6w0H2XD6SMWYsEDa142iVVJVmseFUyc6xCXfa44c7FTABOCo1vr+C/kaKeSuQ2vNpgP5LN5ymOU7j7By91GKSs3CoqiQAFJqhlG3Wgi1o4NJigqmVnhlqlcJvKQzLYvLyjl8opj9x0+x58hJtucUsONwAVn7T3DwRBEAAb4+tEgIp12dKC6rF0XT2HCX/IYSnmHv0ZP8sC2XpdtzWbojl+MnzZYXIZX8SI6pQt3qIdSrFkJSVDA1wytTM7wyIZewOris3MaRwhIOnSgiMSr4kmdVObKQdwB+ADYAvy4tfEJrPe9sXyOF3HWVlNnI3J9H5r481mfnkbX/BLtyC34r7r8Kq+xPeJA/VQL9CQ30I8DPBz8fH/x9FeU2TWm5jZJyGwVFZeSdKiXvVCnHTv55X5gAPx9qRwXToEYoTWLDaRobRkrNMCoHyMHHwvlsNs3O3ALW7s1j7d5jbD6Qz7bDBeSd+vP7NjjAl/CgAMKDfn3v+xLg+/t7v9ymKSm3kV9URn5RKSeKyjhSUMyve9x9PKwlHRtUu6SMsiBIXDKbTbM/7xS7c09yIO8Uh04UcehEMXmnSn97o5aW2ygt15SV2/D1UQT4+eDv60NIJT+qVPYnrLIf0SGBxIQHEhMWSHzVIGIjgvCV1rZwYVprcgtK+OVoIfuOF3Hg+CkOnSjm+KkS8k6Wkl9URnG5jZIy22/vfT9fhZ+PD6GBfoQG+lEl0J9qoZWoViWQ6lUCaREffskHfJ+tkMsOQuK8fHwUsRGm8ArhTZRSRIdWIjq0EmkJVqc5O9k0Swgh3JwUciGEcHNSyIUQws1JIRdCCDcnhVwIIdycFHIhhHBzUsiFEMLNSSEXQgg3J4VcCCHcnBRyIYRwc1LIhRDCzcleK0I4QnEBHFgLx/dC1SSIqg9BchydcAwp5ELYS8lJWPkBrJsKOZtB/3nrX6rUgssfhhY3g49s1SvsRwq5EBVVVgKrJ8CSV6HgICR0gCsehVppEJ4Ax3ZB7lbYPBe+egBWfQjXvASJ7a1OLjyE7EcuREUc2w3TboID6yC+HXR6GhLanvmxWkPWLPjmX3AiG659FVrd7tS4wr3JfuRC2NvmefDFHeb3Az+BRj3hXGecKgWp/aD+NTD9Vpj3MPgFQosbnZNXeCyZtSLExdIaFr8IUwdDRCKMXALJvc5dxP8oIAgGfAx1roI598D6zx2ZVngBKeRCXAybDb5+FL7/DzQbArd+Y4r5xfIPhEGTIbEDzBoJu5bYParwHlLIhbhQ5WUw+y5Y8T60vRt6v2sK8qUKCILBU80Pgtl3mSmLQlwCKeRCXIjyMpg5HNZNgSufhK7PX3hXyrlUCoE+75n55otGVfz5hFeSQi7E+dhs8OV9ZsZJl+fgikfsU8R/Fd8G2vwDVo6TLhZxSaSQC3EuWsOCx2HtJLjiMWh/r2Ouc9VTULU2zL5buljERZNCLsS5fPci/DzGtJg7Pua46wQEQe/34PgeWPqm464jPJIUciHOZuV4+P4laDYUrv63fbtTziShLST3huVj4ORRx15LeBQp5EKcyea5ZsFOvauh51uOL+K/6vg4lBTA0reccz3hEexSyJVSHyqlDiulMu3xfEJYau8Ks/KyZnMY8BH4OnEBdLVGkPidi/oAABepSURBVNofVoyFghznXVe4NXu1yD8GrrHTcwlhndztMGUQVKkJN0yDgGDnZ+j4GJQVSV+5uGB2KeRa6yWAdOoJ91aQA5P7m26UIdMhOMqaHFH1oMkgMx3xxAFrMgi34rQ+cqXUCKXUKqXUqpwcuWUULqbkJHw6CPIPmZZ4ZB1r81zxCJSXwvL3rM0h3ILTCrnWeqzWOl1rnR4dHe2sywpxfrZymDEc9q2G/uMg9m+7hDpf1drQsDusnmh+yAhxDjJrRXg3reHrR2DLXOj2EjTqYXWi37UaAUXHIXOG1UmEi5NCLrzb0jdNX3S7e6H1SKvT/FliB6iWbDbpsuAAGOE+7DX98FPgJ6CBUipbKXWbPZ5XCIdaPw0WjYbU66DzM1an+TulzAlCBzeYKZFCnIW9Zq0M1lrHaK39tdaxWuvx9nheIRxm2yL44k5IvMzsPujjojenjQdCpTBzqLMQZ+Gi714hHGjvCph2o+m2uH4y+FWyOtHZVQqB5kMg6wszo0aIM5BCLrzLoY0weQCExsDQmRAYZnWi82s5HGylZgaLEGfgXoX8wHpY+6nVKYS7OroTJvUD/8pw4ywIcZNpsJF1TBfQuiky6CnOyL0K+eoJ8NX9UHTC6iTC3RzbAxN6QVmxaYlHJFid6OI0u8H8IJJBT3EG7lXImwwye1Bs+tLqJMKd5GXDhJ5QnA83zYbqyVYnuniNeoJ/kGmVC/EX7lXIY1tCRBKs/8zqJMJd5O0zRfzUMdOdEtPE6kSXplIoNOoFmbOgtMjqNMLFuFchV8q0ynctMd+gQpzL0V3w0TVQmAtDZ0CtFlYnqpim10NxHmyZZ3US4WLcq5ADNBkIaMicbnUS4cpytsJH3X7vTolrZXWiiku6HKrUgnVTrU4iXIz7FfLIOhDbCtZJ94o4i/1rTBG3lcMtc92/Jf4rH1/TkNm+CAoOW51GuBD3K+Rg3syHs+CgHEgk/mLbQviou5liOGweVE+xOpF9NR0MutxsLyDEae5ZyFP7g48frJdbTPEHqyea030i68DwReaABk8T3QBimknXovgT9yzkQVWhXlfYMN3cPgvvZiuHhU/DnHugdkfTEg+tYXUqx0ntZ7qPju6yOolwEe5ZyAEaXwf5B+CX5VYnEVY6dQwmX2dOnU+/DW74zEzV82Qpfc2vWbOszSFchvsW8npXg19l2PiF1UmEVQ5mwtgrYdcP0PO/0ON18PW3OpXjhcdDrXQp5OI37lvIK4VAvS6wcbZ0r3gbrWHFB/DBVVB6ynSlpN1sdSrnSu0HB9fDkR1WJxEuwH0LOZhbzIJD8MtPVicRznLyKHw2FOY9bOZV37nUM+aIX6zkPubXzJnW5hAuwb0Lef3T3StZ0r3iFTbPg/fawNYF0PUFc9p9cJTVqawRVgvi2kj3igDcvZAHBEP9rtK94ulOHoWZI2DqYAiuBrf/D9rd7bqn+jhLaj+zniJni9VJhMXc/zshpS8UHoY9y6xOIuzNZjNzw99OMyfJX/EY3P4txDS1OplraNQLUNK9IjygkNfrerp7RW4xPcr+tfDh1WZueHQDGLkErnwc/AKsTuY6qsRAQjtzRyq8mvsX8oBg01e+aY50r3iC43th5kgYe4U5SKHP/8Gwrz1vqb29NOoFOZsgd5vVSYSF3L+QA6T0gcIcmb3izgqPmNWZ75yeH93+frh3tTkZRymr07muRj3Nr9Iq92p+Vgewi7pdwC8QNs6BxA5WpxEX49Rx+OkdWP5/UFJoNkS76imz6EWcX1gtszho0xy4/GGr0wiLeEaLvFII1OlkjoCz2axOIy5EwWFYNBreSIUlr0DdzvCP5dBvrBTxi5XcCw6sg2O7rU4iLOIZhRzMmzl/P+xfbXUScS5HdsDch+DNxvDjm2Z17h0/wsAJUK2h1encU6Ne5lc5y9a1nTxq9gQqLrD7U3tG1wqYAU8fP9NXGJtudRrxR1qb8Yuf3oXNc81+KE0GQvsHIKqu1encX9UkqNHYdC22u8fqNOKvjuwwXYdrJ0PpSYhIhOTedr2EXQq5Uuoa4C3AFxintf6PPZ73olSOgKQrTKuky7MyQOYKSosga6Z5Ex9cb/6PLnsIWo2A0OpWp/MsjXrD4ufhxH6oUtPqNEJr2LkYfn7frET28TONlzb/gBqpdr9chQu5UsoXeBfoAmQDK5VSc7TWGyv63BctuRd8eR8cyjQtFGGNvGxY9RFkfAwncyG6IfR4wxycHRBsdTrPlNzLFPLNc6HV7Van8V7F+eZM1RUfQO4WCI6Gy/8JLW9z6B759miRtwK2a613AiilpgK9AecX8gbd4asHTKtcCrlzaQ27f4QVY00x0TZo0M20vmt3lDskR4tuAFENTNeiFHLny91mivfaKVCSDzWbQ9/3zcpzv0oOv7w9CnktYO8fPs8GWv/1QUqpEcAIgPh4B81KCImG+Hamr/DKJxxzDfFnJYWw/jPzJj680XSftLvbHPIQkWB1Ou+S3At+eA0Kc713MzFnspWbbpMVY003im8ApPQzjZfYNKdGcdpgp9Z6LDAWID09XTvsQsm94OtHzE9ITzyz0VUc222K95pPoCgPajSB3u+a81T9K1udzjs16mWmcm6e6337sztTUR6s/sQU8ON7ILSmWfvQ4mYIqWZJJHsU8n1A3B8+jz39Z9Zo2MMU8o2zZYGEvWkNe1fAT2+bYoEyPzhb3wFxraX7xGo1GpsZEZvmSCF3hKO7zMD9mklQWgjxbaHLM6bmWHwylT0K+UqgnlIqCVPArwdusMPzXpqwWhDbUla62ZPNBlvmmjmw2SshMAza3WtuIcNqWZ1O/Eops2R/+RizYrZyuNWJPEN2Bix7y4y9KV9z19nmTqjZzOpkv6lwIddalyml7gYWYKYffqi1zqpwsopo1AsW/svc/kckWhrFrZWXwobPzcKd3C3m3/LaV6HpYLOaVrieRr1h2dum77bpIKvTuC+tYfcPZsxh53dQ6XTjpfVIl5zeaZc+cq31PGCePZ7LLpJPF/JNX8oCiUtRXgYbpsH3L5kfhtVTof94c7yYr+esIfNItdJMn+2mOVLIL9WuJfDtC7B3uTnIpMuzkH4rVAq1OtlZeeZ3ZUSiGXzbOFsK+cXQ2uw8+O3zcHSH+TccPBXqXyP93+7Cx8d0r6yeYJaCy53Thdu7Er591hTy0Bhz99n8RvAPtDrZeXnOXit/ldzL9OfmWTfu6lb2LINxnWD6MLOT5KDJ5jCHBt2kiLubRj2hrAi2L7Q6iXs4ths+vwXGd4ZDG+HqF+HeNWY+vhsUcfDkQt7o9F4Gm7+yNoery8uGaTfBR93gxAHo/R7c8QM06iEF3F0ltIOgKLOeQpxdcQEsHAXvtIQt8+GKR+G+ddD2H243hdYzu1YAoutDdCPzZm490uo0rqesBJa/C9+/bLpUrnwS2t4NAUFWJxMV5eNrfhCv/xxKT7ldUXI4rU0D7+vH4ES2Gby/6l9uPQPLc1vkYLpXfllm9r4Wv9u32hyltmg01LkK7voZrnhEirgnSelr5jpvk+6VPzmxHz69Hj4baqbR3roA+o5x6yIOHl/Ie5s9PzbJLSZgdiNcOMr0hZ86ZgYyr58sS+k9UUIH070ih5IbWpvNrN5rAzu/h67PmzGg+DZWJ7MLz+1aAaiWbDYSypwFLYdbncZahzfB9FvNfijNh0LXF2TBiCfz9TN3pOumQslJ777bKjwCc+4xi9ri2kCf9yCyjtWp7MqzW+RKQWo/2LPUDOR5I61h1YcwtqM5oPqGz82eKFLEPV9KX3OQwbZvrE5inT3LYEwHM4On6/MwbJ7HFXHw9EIOZjcyNGz8wuokzldcYKZVffUAJLSHO5dB/a5WpxLOktDeLGjJmml1Euez2eD7V+Dj7mYK4W0LzZoSH1+rkzmE5xfy6PpQvTFketmb+cgOGNfZjA90fgaGTLdsZzZhER9fM0609RuHnBPpsorzzWDm4udNQ27kEpfaF8URPL+QA6T2hewVcHzv+R/rCbYtgg+uhIJDcOMs6HC/WfEnvE9KXyg7BdsWWJ3EOY7uhHFdYOt86PYy9B/n0kvr7cU7vrtT+plfvWEEf+V4mDIAwuNhxHfmdB7hveLbQEgN73jv71kGY6+EgoOmAdN6pNcsavOOQl41yRy9lDnD6iSOY7OZqYVzH4R6XWHYfJlWKEz3Skof071y6rjVaRwn6wuY2Md0H96+GGpfYXUip/KOQg6mVX5grek79jTlpTBrBCx90xyxNmiybJYkftdkIJQXm03kPNHyMWZQv2Yzs8CnapLViZzOewp5aj9AwfppViexr9JTZmBnw+fQaRR0f022mhV/VrMFRNYzZ6t6Eq1h8b9h/qPQsDvcNBuCqlqdyhLeU8jDYs3t1rpPTTeEJygugMkDzCEC3V+Hyx70mj5BcRGUgqbXm/UUx/ZYncY+tDZnDnz/klngNnCiV+8p4z2FHKDpDeaw1F+WWZ2k4opOwCd9zQBP3/eh5W1WJxKurMlA86sn3JHabDDvn+YkpJa3Q8+3PXZ++IXyrkLeqAcEhMLaT61OUjG/tsT3r4YBH8lJMOL8wuMh8TJzR6q11Wkundbw9T9h5Qdmgc+1r8jUWrytkAcEQ0pvMxXLXRdIlBTClIHm0Iz+482CDyEuRJNB5uSnfRlWJ7k0WsOCJ2HlOHN+ZpfnpCvxNO8q5ADNhpjtPTd9aXWSi1daZLbg/OUn6DfWTCsT4kIl9zanP61zwztSreF/z5o99FvfYc7RlCL+G+8r5PFtzZme66ZYneTilJeZY9h2LYE+/weNr7M6kXA3gVXM7I7MGaZR4E5+fN18pN0C1/xHivhfeF8hV8oMeu5aAsd/sTrNhbHZTm/DOQ+6vWJmIAhxKVrcZPaid6dN5DI+Nq3xxgOg+xtSxM/A+wo5nC6ECjImWJ3k/H6dZrVuCnR8HFqPsDqRcGdJV5g55Ss+sDrJhdk4x+zeWbeLuROVgc0z8s5/lYgEqH+N+Unv6reYy96Gn96BViPN4bBCVIRS5pCVfatg/xqr05zbriUw4zaolQ4DJ4Cvv9WJXJZ3FnIwG+qczHXtvZrXf25a4yl9pV9Q2E/T68E/yGyw5qoOZcHUIVC1NtzwmZlxJs7Kewt57Y7mGLif33fNebU7v4Mv7jRzf/u+L7eUwn4qh5v+5g3TTX+5q8nbB5OuM8V7yHSvXXZ/MSpUHZRSA5RSWUopm1Iq3V6hnEIp0yo/sBb2rrA6zZ8dzISpQyGqHgyaBH6VrE4kPE3L4Waf8rUuNnvr1HGYfJ05HGLI5xAeZ3Uit1DRZl4m0A9YYocsztf0eqgUBj+PsTrJ7/L2mVWblUJNa0TO1hSOENME4lqbxTWusvdQWYnZAC53K1w/CWo0tjqR26hQIddab9Jab7FXGKcLCIYWN5rtPU/stzoNFOWZIv5raySsltWJhCdrNcKcqLNpjtVJTPfmnLth9w/Q6x05EOUiOa3jVSk1Qim1Sim1Kicnx1mXPb9WtwMafnrX2hxlJTDtJsjdAoMmQo1Ua/MIz5fS10xF/P4l61vli18w2+xe+RQ0G2xtFjd03kKulFqklMo8w8dFbfKhtR6rtU7XWqdHR0dfemJ7i0iEJtebW0yrWuVaw5f3mgHOnv+FOldZk0N4Fx9fM6X18EZrW+UZE2DJK2ax0uUPW5fDjZ23kGutO2utU8/w4TnHjXR8FGzl5s1khW+fM/tfXPkkNB9iTQbhnVL7Wdsq3zL/9IKfzmZPfZlie0lkThuYVnnazbB6oukzdKaV4+GH16DFzXD5P517bSH+2Crf7OSN5PauNEe0xTSBAbLgpyIqOv2wr1IqG2gLzFVKLbBPLAtc/k/w8Yfv/uO8a276EuY9bFaZSmtEWOXXVvl3TmyV524z2zGH1oAbPpczZiuoorNWZmmtY7XWlbTW1bXWV9srmNOF1jD7mKyfBoc2Ov56O76F6bdCrTS47kM5Z1NYx8cXOj4Gh7Mg40PHX+/4XnO6lfKBG2dCiAuNmbkp6Vr5o/b3Q6UqMPch02fuKHtXmOXHUfXNNENZfiysltofal8JC0ebtQyOkn8QJvYyRxXeONMswRcVJoX8j4KqQreXzJmejpqOeGCdWbkWWgNunAWVIxxzHSEuhlLQ803Q5TD3QcdsW3HyKEzsA/mHTAMmpqn9r+GlpJD/VdProWEPM5PE3l0se1fChJ6m1X/TbAipZt/nF6IiIhLhqqdg63xz+IQ9FeSYlvjRnTD4U4hvbd/n93JSyP9KKej5lim2s0aahTr2sPtH+KQPVK4Kw+aZw3CFcDWt7zDjNl8/CicO2Oc5j/8CH14Nudvh+ilQ+wr7PK/4jRTyMwmOgl7/hYPr4ZsnK36buXkuTOoPYbFw63wp4sJ1+fhC73ehrAgm9av47oiHN8P4q82W0Td9AfU62yen+BMp5GfTsDu0vRtWjIX5j11aMbeVmyOqpt4A1ZLhlrmmb1wIV1atEVw/GY5shymDoKTw0p4n6wsY3xVsZXDLPIhvY9+c4jdSyM+l6/PQ5i6zO+K8f15cMS/MNYOaP7xmlh4P+9q09IVwB7U7Qv/xkL3S7AFUeurCv7a0yMz8+vxmiKoLwxfJ3kEOJpOXz0UpuPoFc6jDsrch/wB0eRYi65z9a0pPwfL34Mc3ze1pz7fMyd9CuJvkXtDjTbMP0HttzKK1up3O/nibzezZ8t2LkLPZ3NF2GgV+Ac7L7KWkkJ+PUtDlOQiKhO9fNiP6LW6G9FvN2Z+VQqGsGPatNgOaGR/BiX1Qvxt0Hg3VGlr9CoS4dGk3Q9Uksx/KpH6Qep2Z2RXTzCzkKS81qzT3rTJTdnM2Q2Rds1qzfler03sNpS045iw9PV2vWrXK6detsPxDZnOh1RNMvx9AYLhpeZedPsQ5rjV0ehoSO1iXUwh7Ky2CH9+AH1+H8tMzuYKrmcFQW6n5PLqh2eoipa8ZNBV2p5TK0Fr/7TQ2KeSX4tge03eYt9csN/YLhIR25kPOFxSerOiEWdR2YK1ZZxFSDaqnmMH8aslytqyDna2QS9fKpYhIMB9CeJvAKpB0mfkQLkN+fAohhJuTQi6EEG5OCrkQQrg5KeRCCOHmpJALIYSbk0IuhBBuTgq5EEK4OSnkQgjh5ixZ2amUygH2OP3CFRcF5Fodwom87fWCvGZv4a6vOUFr/bfTqi0p5O5KKbXqTMtjPZW3vV6Q1+wtPO01S9eKEEK4OSnkQgjh5qSQX5yxVgdwMm97vSCv2Vt41GuWPnIhhHBz0iIXQgg3J4VcCCHcnBTyS6CUekgppZVSUVZncTSl1CtKqc1KqfVKqVlKqXCrMzmKUuoapdQWpdR2pdRjVudxNKVUnFJqsVJqo1IqSyl1n9WZnEEp5auUWqOU+srqLPYihfwiKaXigK7AL1ZncZKFQKrWugmwFXjc4jwOoZTyBd4FugHJwGClVLK1qRyuDHhIa50MtAHu8oLXDHAfsMnqEPYkhfzivQE8AnjFKLHW+hut9emTplkOxFqZx4FaAdu11ju11iXAVKC3xZkcSmt9QGu9+vTv8zHFrZa1qRxLKRULdAfGWZ3FnqSQXwSlVG9gn9Z6ndVZLHIr8LXVIRykFrD3D59n4+FF7Y+UUolAc+Bna5M43JuYhpjN6iD2JIcv/4VSahFQ4wx/9STwBKZbxaOc6zVrrWeffsyTmFvxyc7MJhxPKRUCzADu11qfsDqPoyilegCHtdYZSqmOVuexJynkf6G17nymP1dKNQaSgHVKKTBdDKuVUq201gedGNHuzvaaf6WUugXoAXTSnrvwYB8Q94fPY0//mUdTSvljivhkrfVMq/M4WHugl1LqWiAQqKKUmqS1HmpxrgqTBUGXSCm1G0jXWrvjDmoXTCl1DfA6cIXWOsfqPI6ilPLDDOZ2whTwlcANWussS4M5kDItkgnAUa31/VbncabTLfKHtdY9rM5iD9JHLs7nHSAUWKiUWquUGmN1IEc4PaB7N7AAM+g3zZOL+GntgRuBq07/36493VoVbkZa5EII4eakRS6EEG5OCrkQQrg5KeRCCOHmpJALIYSbk0IuhBBuTgq5EEK4OSnkQgjh5v4fulcbD9k6UPgAAAAASUVORK5CYII=\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "x_grid = jnp.linspace(-5, 5, 100)\n",
        "value, grad = jax.vmap(jax.value_and_grad(jitted_function))(x_grid)\n",
        "plt.plot(x_grid, value, label=\"value\")\n",
        "plt.plot(x_grid, grad, label=\"grad\")\n",
        "plt.legend();"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f0704757-2bd8-4439-a0f7-8d17dd1a138a",
      "metadata": {
        "id": "f0704757-2bd8-4439-a0f7-8d17dd1a138a"
      },
      "source": [
        "## PyTrees\n",
        "\n",
        "Another useful JAX concept is \"PyTrees\".\n",
        "This allows us to use structured inputs and still use `jit`, `vmap`, and `grad`.\n",
        "For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "db182d2d-47c3-4f63-88cc-f0f316fb0ad8",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "db182d2d-47c3-4f63-88cc-f0f316fb0ad8",
        "outputId": "32cf977c-757a-42fa-d6c7-af71c435c4fa"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(0.02227585, dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ],
      "source": [
        "def pytree_func(params):\n",
        "    return jnp.exp(params[\"log_amp\"]) * jnp.sin(params[\"log_scale\"])\n",
        "\n",
        "params = {\n",
        "    \"log_amp\": -1.5,\n",
        "    \"log_scale\": 0.1,\n",
        "}\n",
        "pytree_func(params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "id": "59f7514e-4c74-4a64-9f56-bfec4469c41c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "59f7514e-4c74-4a64-9f56-bfec4469c41c",
        "outputId": "8b8e8341-ed82-487f-9387-da348e724b9a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'log_amp': DeviceArray(0.02227585, dtype=float32, weak_type=True),\n",
              " 'log_scale': DeviceArray(0.22201544, dtype=float32, weak_type=True)}"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ],
      "source": [
        "jax.grad(pytree_func)(params)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Random numbers\n",
        "\n",
        "Random number generation in JAX is a little different from in numpy.\n",
        "For example, every random function takes a \"key\" as input:"
      ],
      "metadata": {
        "id": "RFR4UqrmMoFQ"
      },
      "id": "RFR4UqrmMoFQ"
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import random\n",
        "\n",
        "key = random.PRNGKey(42)\n",
        "random.normal(key)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cTAI87MPMnvI",
        "outputId": "8c93bd9c-31ab-4cfe-abb1-f5dd7cd04135"
      },
      "id": "cTAI87MPMnvI",
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(-0.18471177, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "id": "0b6d64a9-6361-4f05-b399-29312c15ef31",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0b6d64a9-6361-4f05-b399-29312c15ef31",
        "outputId": "8644972a-d897-4877-c0be-e1dd64da92ef"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(-0.18471177, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "random.normal(key)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you want to generate multiple different random numbers, a good approach is to \"split\" the key."
      ],
      "metadata": {
        "id": "sOY38c39NQNN"
      },
      "id": "sOY38c39NQNN"
    },
    {
      "cell_type": "code",
      "source": [
        "key1, key2 = random.split(key)\n",
        "random.normal(key1), random.uniform(key2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9v5Hob9_NEKM",
        "outputId": "06b28e51-f6ed-40a5-8845-f95ad1a81292"
      },
      "id": "9v5Hob9_NEKM",
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(0.13790321, dtype=float32),\n",
              " DeviceArray(0.91457367, dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optimizers\n",
        "\n",
        "The JAX ecosystem is pretty modular and there are various packages available for non-linear function optimization.\n",
        "Some popular ones include [jaxopt](https://github.com/google/jaxopt) (\"scipy.optimize with support for PyTrees\") and [optax](https://github.com/deepmind/optax) (\"feature-rich framework with a lot more boilerplate\")."
      ],
      "metadata": {
        "id": "u01iPar9Wc8B"
      },
      "id": "u01iPar9Wc8B"
    },
    {
      "cell_type": "code",
      "source": [
        "%pip install -q jaxopt optax"
      ],
      "metadata": {
        "id": "LPVdUpTUNLLU"
      },
      "id": "LPVdUpTUNLLU",
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jaxopt\n",
        "import optax\n",
        "\n",
        "def loss(params):\n",
        "  return jnp.sum(jnp.square(params[\"x\"]))\n",
        "\n",
        "params = {\"x\": 12.5}\n",
        "opt = jaxopt.ScipyMinimize(fun=loss)\n",
        "soln = opt.run(params)\n",
        "print(soln)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1oYP60kxXf1v",
        "outputId": "e67eddb6-d155-4d25-f599-4607aef8d41e"
      },
      "id": "1oYP60kxXf1v",
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "OptStep(params={'x': DeviceArray(4.7211597e-07, dtype=float32)}, state=ScipyMinimizeInfo(fun_val=DeviceArray(2.228935e-13, dtype=float32, weak_type=True), success=True, status=0, iter_num=2))\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "params = {\"x\": 12.5}\n",
        "opt = optax.sgd(0.1)\n",
        "opt_state = opt.init(params)\n",
        "\n",
        "@jax.jit\n",
        "def train(params, opt_state):\n",
        "  value, grads = jax.value_and_grad(loss)(params)\n",
        "  updates, opt_state = opt.update(grads, opt_state)\n",
        "  params = optax.apply_updates(params, updates)\n",
        "  return value, params, opt_state\n",
        "\n",
        "losses = []\n",
        "for _ in range(100):\n",
        "  value, params, opt_state = train(params, opt_state)\n",
        "  losses.append(value)\n",
        "params"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4FUgjJwXXuq8",
        "outputId": "bf12e9a8-d442-414d-f23a-cb0693cca0b7"
      },
      "id": "4FUgjJwXXuq8",
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'x': DeviceArray(2.5462958e-09, dtype=float32)}"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(losses)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 282
        },
        "id": "rF30CR93YdEP",
        "outputId": "9c201c4c-163a-4b95-b542-e57bdf4e9aa4"
      },
      "id": "rF30CR93YdEP",
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x7fbb7f52a710>]"
            ]
          },
          "metadata": {},
          "execution_count": 28
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD4CAYAAAAXUaZHAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAY1UlEQVR4nO3da5Bc5X3n8e9vZnpukmZGl9EIJBEJGKCw115TAwjsuDDYMfgmXrhiSDaWs0qpnGUTJ/auA+uqZfeFq+xdV4i92aVWCxjhpcAEE0O5iB0sY1PZGOERdgCJiwZhYFSSZkDoLjS3/77o06g1mmFGfZlWn/P7VKmm+5zT3f/jg3/99NNPP48iAjMzS5eGWhdgZmaV53A3M0shh7uZWQo53M3MUsjhbmaWQk21LgBgyZIlsWrVqlqXYWZWV7Zu3fpGRHRPte+MCPdVq1bR399f6zLMzOqKpFen2zdjt4ykuyQNSXpu0vY/k/SCpG2S/lvR9lskDUh6UdLHyyvdzMxKMZuW+93A3wL3FDZI+giwFnh/RByXtDTZfjFwA/Ae4Gzgp5IuiIjxShduZmbTm7HlHhFPAPsmbf5T4BsRcTw5ZijZvha4PyKOR8QrwABwWQXrNTOzWSh1tMwFwO9K2iLpF5IuTbYvB14vOm4w2XYKSRsk9UvqHx4eLrEMMzObSqnh3gQsAtYA/xF4QJJO5wkiYmNE9EVEX3f3lF/2mplZiUoN90Hgoch7CpgAlgC7gJVFx61ItpmZ2RwqNdx/CHwEQNIFQDPwBvAIcIOkFkmrgV7gqUoUamZmszeboZD3Ab8ELpQ0KGk9cBdwbjI88n5gXdKK3wY8AGwHfgzcVM2RMi/uOcS3fvIibx4+Xq2XMDOrSzMOhYyIG6fZ9W+mOf7rwNfLKWq2dg4f5m8fH+CT7zuLxfNb5uIlzczqQl3PLdPW3AjA0ZGxGldiZnZmqetwb2/Of/A4OuLfSJmZFavzcC+03B3uZmbF6jrcC90yxxzuZmYnqetwd8vdzGxq9R3uuUKfu79QNTMrVtfh7m4ZM7Op1XW4Nzc10NQgjo463M3MitV1uEO+9e6Wu5nZyeo+3NubG93nbmY2SQrCvcmjZczMJqn7cG/LuVvGzGyyug/3fLeMw93MrFj9h3tLk0fLmJlNUv/hnmvkmL9QNTM7Sf2Hu7tlzMxOUffh7nHuZmanms0ye3dJGkqW1Ju87yuSQtKS5L4kfUfSgKRnJF1SjaKLueVuZnaq2bTc7waunbxR0krg94DXijZfR35R7F5gA3B7+SW+u7bmJo6NjjMxEdV+KTOzujFjuEfEE8C+KXbdBnwVKE7VtcA9yWLZTwJdks6qSKXTKEz7e8wjZszM3lFSn7uktcCuiPiXSbuWA68X3R9Mtk31HBsk9UvqHx4eLqUMwHO6m5lN5bTDXVI78J+A/1zOC0fExojoi4i+7u7ukp+nLedpf83MJmsq4THnAauBf5EEsAJ4WtJlwC5gZdGxK5JtVfPOItmjHutuZlZw2i33iHg2IpZGxKqIWEW+6+WSiNgDPAJ8Phk1swY4EBG7K1vyydwtY2Z2qtkMhbwP+CVwoaRBSevf5fBHgZ3AAPB/gH9XkSrfhVdjMjM71YzdMhFx4wz7VxXdDuCm8suaPbfczcxOVfe/UD0R7u5zNzMrqPtwb0u+UHW3jJnZCXUf7u05d8uYmU1W9+He5l+ompmdou7DvaWpgQa5z93MrFjdh7skL5JtZjZJ3Yc75EfM+AtVM7MTUhPubrmbmZ2QinBvc7eMmdlJUhHu+Za7v1A1MytIUbi75W5mVpCKcG/L+QtVM7NiqQj39uZGz+duZlYkFeHe1tzklruZWZFUhLv73M3MTpaacD82Ok5+OnkzM5vNSkx3SRqS9FzRtv8u6QVJz0j6e0ldRftukTQg6UVJH69W4cXamhuJgLdHJ+bi5czMznizabnfDVw7adtjwHsj4n3AS8AtAJIuBm4A3pM85n9JaqxYtdM4Me2vv1Q1M4NZhHtEPAHsm7TtHyOikKRPAiuS22uB+yPieES8Qn4t1csqWO+U2pMFO9zvbmaWV4k+938L/ENyeznwetG+wWTbKSRtkNQvqX94eLisAjynu5nZycoKd0lfA8aAe0/3sRGxMSL6IqKvu7u7nDK8SLaZ2SRNpT5Q0heATwHXxIlhKruAlUWHrUi2VVWbF8k2MztJSS13SdcCXwU+ExFHi3Y9AtwgqUXSaqAXeKr8Mt9duxfJNjM7yYwtd0n3AVcBSyQNAreSHx3TAjwmCeDJiPhiRGyT9ACwnXx3zU0RUfXEdbeMmdnJZgz3iLhxis13vsvxXwe+Xk5Rp6sQ7m65m5nlpeQXqoWhkO5zNzOD1IR70i3joZBmZkBKwr2lqQEJjh53uJuZQUrCXRLtOc8MaWZWkIpwh2ROdy/YYWYGpCjcPae7mdkJDnczsxRKTbi3NXuRbDOzgtSEe77l7j53MzNIUbi35ZrcLWNmlkhNuBfWUTUzs5SFu1vuZmZ5qQl3f6FqZnZCasK98IXqiXVDzMyyK0Xh3sREwPGxiVqXYmZWc6kJ97ac53Q3MyuYMdwl3SVpSNJzRdsWSXpM0o7k78JkuyR9R9KApGckXVLN4ot52l8zsxNm03K/G7h20rabgc0R0QtsTu4DXEd+3dReYANwe2XKnFnbO6sx+YdMZmYzhntEPAHsm7R5LbApub0JuL5o+z2R9yTQJemsShX7bua9sxqTW+5mZqX2ufdExO7k9h6gJ7m9HHi96LjBZNspJG2Q1C+pf3h4uMQyTmhvybfcDx93y93MrOwvVCM/9vC0xx9GxMaI6IuIvu7u7nLLoKM1B8DBYw53M7NSw31vobsl+TuUbN8FrCw6bkWyreq62vPhfuDYyFy8nJnZGa3UcH8EWJfcXgc8XLT988momTXAgaLum6rqbCuE++hcvJyZ2RmtaaYDJN0HXAUskTQI3Ap8A3hA0nrgVeD3k8MfBT4BDABHgT+uQs1Tmt/SRGODHO5mZswi3CPixml2XTPFsQHcVG5RpZBER2sT+4863M3MUvMLVYCu9ma33M3MSFm4d7TlHO5mZqQs3Dvbchx0uJuZpSvcu9py7He4m5mlK9w73S1jZgakMNwPHhtlYsILdphZtqUq3Lvac0wEHPbMkGaWcakK947Cr1Q91t3MMi5V4e4pCMzM8hzuZmYplKpwPzEzpMPdzLItVeFeaLl7fhkzy7pUhrtb7maWdakK97ZcI82NDQ53M8u8VIW7pGTyMK/GZGbZlqpwB+hsa3LL3cwyL4Xh7vllzMzKCndJfylpm6TnJN0nqVXSaklbJA1I+r6k5koVOxtesMPMrIxwl7Qc+HOgLyLeCzQCNwDfBG6LiPOBt4D1lSh0tjrbch4KaWaZV263TBPQJqkJaAd2A1cDDyb7NwHXl/kap8XdMmZmZYR7ROwCvgW8Rj7UDwBbgf0RUZiWcRBYPtXjJW2Q1C+pf3h4uNQyTtHZluPQ22OMe9pfM8uwcrplFgJrgdXA2cA84NrZPj4iNkZEX0T0dXd3l1rGKQo/ZDr0tlvvZpZd5XTLfBR4JSKGI2IUeAj4INCVdNMArAB2lVnjafEUBGZm5YX7a8AaSe2SBFwDbAceBz6bHLMOeLi8Ek+PJw8zMyuvz30L+S9OnwaeTZ5rI/BXwJclDQCLgTsrUOeseX4ZM7P8aJeSRcStwK2TNu8ELivnecvhcDczS+kvVAH2O9zNLMNSF+6FdVQPOtzNLMNSF+6tuUZac57218yyLXXhDoUpCDztr5llVyrDvavNk4eZWbalMtw9v4yZZV0qwz2/GtPYzAeamaVUKsO9sy3HAfe5m1mGpTLcu9rdLWNm2ZbKcO9sy3FkZJzR8Ylal2JmVhOpDXfwD5nMLLtSGe6FmSE9BYGZZVUqw73Dk4eZWcalMtzfmRnSC3aYWUalMtyXzGsBYPjw8RpXYmZWG6kM96Ud+XDfe+DtGldiZlYbZYW7pC5JD0p6QdLzkq6QtEjSY5J2JH8XVqrY2WrNNbKwPceegw53M8umclvu3wZ+HBEXAe8HngduBjZHRC+wObk/53o6WtnrcDezjCo53CV1Ah8mWSM1IkYiYj+wFtiUHLYJuL7cIkuxrLPVLXczy6xyWu6rgWHgu5J+LekOSfOAnojYnRyzB+iZ6sGSNkjql9Q/PDxcRhlTW9bRyp4D/kLVzLKpnHBvAi4Bbo+IDwBHmNQFExEBxFQPjoiNEdEXEX3d3d1llDG1no5W3jxy3FMQmFkmlRPug8BgRGxJ7j9IPuz3SjoLIPk7VF6JpVnW2UoEDB1y693MsqfkcI+IPcDrki5MNl0DbAceAdYl29YBD5dVYYmWdbQCsMfDIc0sg5rKfPyfAfdKagZ2An9M/g3jAUnrgVeB3y/zNUrSk4S7R8yYWRaVFe4R8Rugb4pd15TzvJWwrNMtdzPLrlT+QhVgYXuO5qYGt9zNLJNSG+6S6Olo8Vh3M8uk1IY7FMa6O9zNLHtSHe6egsDMsirV4b6sIz8FQf63VGZm2ZHucO9s5e3RCQ4eG6t1KWZmcyrV4V4Y6+4vVc0sa1Id7u+MdXe4m1nGpDvcC79S9YgZM8uYVId7Ybk9t9zNLGtSHe4tTY0smtfscDezzEl1uEMy1t3dMmaWMakP92WegsDMMij94d7pX6maWfakPtx7Olp54/AII2Nebs/MsiP14V4YDjl0yK13M8uOssNdUqOkX0v6UXJ/taQtkgYkfT9Zpalmerxoh5llUCVa7l8Cni+6/03gtog4H3gLWF+B1yjZMk9BYGYZVFa4S1oBfBK4I7kv4GrgweSQTcD15bxGuVYsbAPgtX1Ha1mGmdmcKrfl/jfAV4HCt5WLgf0RUZiGcRBYPtUDJW2Q1C+pf3h4uMwypregNceyjlYG9h6u2muYmZ1pSg53SZ8ChiJiaymPj4iNEdEXEX3d3d2lljErvT3z2THkcDez7Cin5f5B4DOSfgvcT7475ttAl6Sm5JgVwK6yKqyA3qULGBg6zMSEF+0ws2woOdwj4paIWBERq4AbgJ9FxB8CjwOfTQ5bBzxcdpVl6u2Zz7HRcXbtP1brUszM5kQ1xrn/FfBlSQPk++DvrMJrnJbepfMB2DF0qMaVmJnNjaaZD5lZRPwc+HlyeydwWSWet1LOL4T73sNcfVFPjasxM6u+1P9CFaCrvZnuBS3+UtXMMiMT4Q5wgUfMmFmGZCbce5cuYGDvISI8YsbM0i8z4X7+0vkcGRlnt+eYMbMMyEy4nxgx464ZM0u/7IR7zwIAduz1cEgzS7/MhPuiec0smd/MDs8xY2YZkJlwh3y/u3/IZGZZkKlw7126gB1Dhz1ixsxSL1vh3jOfQ2+PMXToeK1LMTOrqkyFe/E0BGZmaZapcL8gGTHzwp6DNa7EzKy6MhXuS+a3sHJRG7/67b5al2JmVlWZCneAK85dzJM79zHuhTvMLMUyF+5XnreEA8dGeX63u2bMLL0yF+5XnLcYgH9++Y0aV2JmVj3lLJC9UtLjkrZL2ibpS8n2RZIek7Qj+buwcuWWr6ejlfO65/HPL79Z61LMzKqmnJb7GPCViLgYWAPcJOli4GZgc0T0ApuT+2eUK89bwlOv7GN0fKLWpZiZVUU5C2Tvjoink9uHgOeB5cBaYFNy2Cbg+nKLrLQrz1vM0ZFxnhncX+tSzMyqoiJ97pJWAR8AtgA9EbE72bUHOOMWLV1zbtLvPuCuGTNLp7LDXdJ84AfAX0TESUNQIj+Jy5RjDiVtkNQvqX94eLjcMk7LwnnNXHxWh/vdzSy1ygp3STnywX5vRDyUbN4r6axk/1nA0FSPjYiNEdEXEX3d3d3llFGSK89bzNbX3uLt0fE5f20zs2orZ7SMgDuB5yPir4t2PQKsS26vAx4uvbzqufL8xYyMTfD0a2/VuhQzs4orp+X+QeCPgKsl/Sb59wngG8DHJO0APprcP+NcumoRjQ3iFy/NbZeQmdlcaCr1gRHxT4Cm2X1Nqc87Vxa05rjqgm4eenoX/+H3LiTXmLnfc5lZimU60f7g8nMYPnScn27fW+tSzMwqKtPhftWFSzm7s5V7t7xW61LMzCoq0+He2CA+d+k5/NPAG/z2jSO1LsfMrGIyHe4An7t0JY0N4r5fufVuZumR+XBf1tnK1Rct5cH+QUbGPNeMmaVD5sMd8l+svnlkhJ9s21PrUszMKsLhDny4t5vfWdzOdzbvcOvdzFLB4U7+i9X/8un3sGPoMP/7Fy/Xuhwzs7I53BMfuWgpn3zfWfyPxwfYOXy41uWYmZXF4V7k1k9fTEtTA1/7++fIT2hpZlafHO5Fli5o5ebrLuKXO9/k7/oHa12OmVnJHO6T3HjpOVy+ehFf++Gz/OwFT0tgZvXJ4T5JQ4PY+Ed9XLhsAV/83tOeNdLM6pLDfQqd7Tn+7/rLOW/pfDbc0++AN7O643CfRld7M/f+yeWsXjKPL3z3KW59+DkOHx+rdVlmZrPicH8Xi+Y184M/vZIvXLmKe558lY/f9gT/uG0PExMeSWNmZzaH+wzmtTRx66ffw4NfvILWXAMbvreVq771czY+8TJvHRmpdXlmZlNStcZzS7oW+DbQCNwREdMut9fX1xf9/f1VqaOSRsYm+Mm2PXzvyVd56pV9NAj+1YouPnT+Ytacu5iLlnWwZH4z+eVlzcyqS9LWiOibcl81wl1SI/AS8DFgEPgVcGNEbJ/q+HoJ92Iv7DnIo8/u4f8NvMFvXt/PeNJVs7A9x/lL53N2VxvLOlvpWdDKwnk5utqa6WjLMb+lifbmRtqaG2lpaqClqZFco/yGYGan7d3CveQ1VGdwGTAQETuTAu4H1gJThns9umhZBxct6+DLH7uAg2+P8szrB3hp7yFe2nuIncNHePq1t9h74Dgj47ObiCzXKJoaGmhqEA0NorFBNEg0NkCD8rcBpPz9wnuB4J03hnfeHoreJ6Z7yyjnzcRvQ2aV87lLV/Inv3tuxZ+3WuG+HHi96P4gcHnxAZI2ABsAzjnnnCqVMTc6WnN8qHcJH+pdctL2iYlg/7FR9h8d4cCxUfYfG+XYyDhHjo9xdGSckbEJRsYnOD42wdj4BGMTwej4BBMTwXgE4xMQEUwUbhMQMJF82gqg8MGr8Pmr+JPYtJ/JyviwFuU82MxOsWR+S1Wet1rhPqOI2AhshHy3TK3qqKaGBrFoXjOL5jXXuhQzy5hqjZbZBawsur8i2WZmZnOgWuH+K6BX0mpJzcANwCNVei0zM5ukKt0yETEm6d8DPyE/FPKuiNhWjdcyM7NTVa3PPSIeBR6t1vObmdn0/AtVM7MUcribmaWQw93MLIUc7mZmKVS1icNOqwhpGHi1xIcvAd6oYDn1IovnncVzhmyedxbPGU7/vH8nIrqn2nFGhHs5JPVPN3FOmmXxvLN4zpDN887iOUNlz9vdMmZmKeRwNzNLoTSE+8ZaF1AjWTzvLJ4zZPO8s3jOUMHzrvs+dzMzO1UaWu5mZjaJw93MLIXqOtwlXSvpRUkDkm6udT3VIGmlpMclbZe0TdKXku2LJD0maUfyd2Gta60GSY2Sfi3pR8n91ZK2JNf8+8mU0qkhqUvSg5JekPS8pCuycK0l/WXy3/dzku6T1JrGay3pLklDkp4r2jbl9VXed5Lzf0bSJafzWnUb7ski3P8TuA64GLhR0sW1raoqxoCvRMTFwBrgpuQ8bwY2R0QvsDm5n0ZfAp4vuv9N4LaIOB94C1hfk6qq59vAjyPiIuD95M891dda0nLgz4G+iHgv+WnCbyCd1/pu4NpJ26a7vtcBvcm/DcDtp/NCdRvuFC3CHREjQGER7lSJiN0R8XRy+xD5/7MvJ3+um5LDNgHX16bC6pG0AvgkcEdyX8DVwIPJIak6b0mdwIeBOwEiYiQi9pOBa01++vE2SU1AO7CbFF7riHgC2Ddp83TXdy1wT+Q9CXRJOmu2r1XP4T7VItzLa1TLnJC0CvgAsAXoiYjdya49QE+NyqqmvwG+Ckwk9xcD+yNiLLmftmu+GhgGvpt0Rd0haR4pv9YRsQv4FvAa+VA/AGwl3de62HTXt6yMq+dwzxRJ84EfAH8REQeL90V+PGuqxrRK+hQwFBFba13LHGoCLgFuj4gPAEeY1AWT0mu9kHwrdTVwNjCPU7suMqGS17eewz0zi3BLypEP9nsj4qFk897CR7Tk71Ct6quSDwKfkfRb8l1uV5Pvj+5KPrpD+q75IDAYEVuS+w+SD/u0X+uPAq9ExHBEjAIPkb/+ab7Wxaa7vmVlXD2HeyYW4U76me8Eno+Ivy7a9QiwLrm9Dnh4rmurpoi4JSJWRMQq8tf2ZxHxh8DjwGeTw1J13hGxB3hd0oXJpmuA7aT8WpPvjlkjqT35771w3qm91pNMd30fAT6fjJpZAxwo6r6ZWUTU7T/gE8BLwMvA12pdT5XO8UPkP6Y9A/wm+fcJ8v3Pm4EdwE+BRbWutYr/G1wF/Ci5fS7wFDAA/B3QUuv6Knyu/xroT673D4GFWbjWwH8FXgCeA74HtKTxWgP3kf9eYZT8J7X1011fQORHBL4MPEt+NNGsX8vTD5iZpVA9d8uYmdk0HO5mZinkcDczSyGHu5lZCjnczcxSyOFuZpZCDnczsxT6/6F/oXCwfCMWAAAAAElFTkSuQmCC\n"
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "Afl22IuwYie5"
      },
      "id": "Afl22IuwYie5",
      "execution_count": 28,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.9"
    },
    "colab": {
      "name": "intro-to-jax-part2.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}