Jax_first.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Jax.first.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rahuldave/LearningJax/blob/main/Jax_first.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SpqDrFxx7vsi"
      },
      "source": [
        "# Jax.first\n",
        "\n",
        "What is Jax? Its readme describes it as: \"AX is really an extensible system for composable function transformations.\", which does not tell us much.\n",
        "\n",
        "A more germane question might be why use it? The answer is, it enables many different kinf=d of interesting applications, all of which run very fast.\n",
        "\n",
        "For example, you can make BERT 12 tmes faster than the default implementation.\n",
        "\n",
        "Yes, neural networks. But also gradient boosting, solving diffferential equations, high performance computing, etc.\n",
        "\n",
        "How?\n",
        "\n",
        "![](https://i.imgur.com/COjvmnE.png)\n",
        "\n",
        "This diagram and a detailed discussion on how why|when you should use jax can be found at https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/#references .\n",
        "\n",
        "Let us address its features one by one."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cat /proc/cpuinfo"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "47_EXkauOvTK",
        "outputId": "6ea12933-219d-4a93-ef4d-38658f2fb25c"
      },
      "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": "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": "6d675079-9d33-476f-f4d0-deeb99e6f4a8"
      },
      "source": [
        "num_replicas = len(jax.devices())\n",
        "jax.devices()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
              " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
              " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
              " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KR_PE58Mhzky"
      },
      "source": [
        "## 1:  `jax.numpy` replaces `numpy` and makes it faster\n",
        "\n",
        "jax.numpy provides a dropin replacement for numpy (*) that runs ops on whatever accelerator you have.\n",
        "\n",
        "This section uses code from https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/jax_intro.ipynb#scrollTo=D2dSgG-cnVIe"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "il2BnQpUh15C"
      },
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Everything (almost) that you might call in numpy can be obtained using jax. The advatage is that there is a bunch of optimizations you can take advatage of..."
      ],
      "metadata": {
        "id": "6xBlCSewmQC5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "size = int(1e3)\n",
        "number_of_loops=int(1e2)"
      ],
      "metadata": {
        "id": "UmPlcaBORtKz"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def f(x=None):\n",
        "  if not isinstance(x, np.ndarray):\n",
        "    x=np.ones((size, size), dtype=np.float32) \n",
        "  return np.dot(x, x.T)"
      ],
      "metadata": {
        "id": "SQ4YXgweRaAS"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit -o -n $number_of_loops f()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2bcyi_t1RcP5",
        "outputId": "881f8d29-dae0-4ffa-9b9a-ac1cf59e08e3"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100 loops, best of 5: 44.7 ms per loop\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<TimeitResult : 100 loops, best of 5: 44.7 ms per loop>"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "JAX supports execution on XLA devices through a process called `jit`ting, in which your python code is compiled down firtst to an intermediate representation language, and finally using XLA into machine code appropriate for a CPU, a GPU, or a TPU."
      ],
      "metadata": {
        "id": "QNAjkElHm3zr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# JAX device execution\n",
        "# https://github.com/google/jax/issues/1598\n",
        "\n",
        "def jf(x=None): \n",
        "  if not isinstance(x, jnp.ndarray):\n",
        "    x=jnp.ones((size, size), dtype=jnp.float32)\n",
        "  return jnp.dot(x, x.T)"
      ],
      "metadata": {
        "id": "bQtUysVJR65z"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import jit\n",
        "f_tpu = jit(jf)\n",
        "f_cpu = jit(jf, backend='cpu')"
      ],
      "metadata": {
        "id": "OhffnpEfSya0"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "  %time f_cpu() "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rVkQ0C0LUHA3",
        "outputId": "c472598f-366e-4659-ad6b-49020df6c5a4"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 224 ms, sys: 28.9 ms, total: 253 ms\n",
            "Wall time: 173 ms\n",
            "CPU times: user 96.9 ms, sys: 10 ms, total: 107 ms\n",
            "Wall time: 58.4 ms\n",
            "CPU times: user 95 ms, sys: 0 ns, total: 95 ms\n",
            "Wall time: 53.5 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Sometimes JAX `jit`ted CPU can be slower than numpy CPU, but if a function is complex, this is usually not the case."
      ],
      "metadata": {
        "id": "YLOSkF6EqOMv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(3):\n",
        "  %time f_tpu() "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-q2KD3_8mBHy",
        "outputId": "05de62c3-07d4-4b67-dff6-8ac1dfcb7056"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 198 ms, sys: 436 ms, total: 633 ms\n",
            "Wall time: 1.11 s\n",
            "CPU times: user 10.5 ms, sys: 10.5 ms, total: 20.9 ms\n",
            "Wall time: 30.7 ms\n",
            "CPU times: user 12.5 ms, sys: 18 ms, total: 30.5 ms\n",
            "Wall time: 49.1 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit -o -n $number_of_loops f_tpu().block_until_ready() "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1aNdG8kiUjdF",
        "outputId": "b43297a9-ce24-45ba-a7c4-898549cbfe0c"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100 loops, best of 5: 1.88 ms per loop\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<TimeitResult : 100 loops, best of 5: 1.88 ms per loop>"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Why is there a `block_until_ready` there? It is needed to get accurate timings. This is because JAX does not:\n",
        "\n",
        ">wait for the operation to complete before returning control to the Python program. Instead, JAX returns a DeviceArray value, which is a future, i.e., a value that will be produced in the future on an accelerator device but isn’t necessarily available immediately. We can inspect the shape or type of a DeviceArray without waiting for the computation that produced it to complete, and we can even pass it to another JAX computation, as we do with the addition operation here. Only if we actually inspect the value of the array from the host, for example by printing it or by converting it into a plain old numpy.ndarray will JAX force the Python code to wait for the computation to complete.\n",
        "\n",
        "(from https://jax.readthedocs.io/en/latest/async_dispatch.html)\n",
        "\n",
        "The `DeviceArray` here replaces numpy's standard `ndarray`. Because of the design above, it is furthermore \"lazy\", which is slightly different from the general concept of a lazy computation.\n",
        "\n",
        "For example computations in Spark are lazy because they can be composed, the idea being that a compiler can figure how to \"fuse\" operations to make them more performant. For example, in spark, if only a portion of a dataframe is required, the compiler might figure that the previous computation didnt need the whole dataframe either and optimize.\n",
        "\n",
        "Pretty much the same happens here, except that the result of a `DeviceArray` computation is kept on the device where it is being done, for example the TPU. Indeed, if you `jit` a function, the XLA compiler might be able to figure a faster code path and run that instead."
      ],
      "metadata": {
        "id": "TRCUF-6DnvfX"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "csCNakq2iPJx"
      },
      "source": [
        "## 2: Function transformation\n",
        "\n",
        "The JAX readme tells us that JAX is a function composition engine, and indeed this is true. Consider a simple quadratic function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TDokNNDZcktr",
        "outputId": "acf0236a-bb5b-46b7-c84c-c266b203ee6b"
      },
      "source": [
        "def f(x):\n",
        "  return 2*x*x + 3*x + 3\n",
        "\n",
        "f(3)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "30"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-lVTCRLgbIj2"
      },
      "source": [
        "we can use `make_jaxpr` to trace the function and show us a jax expression of what the function does"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AGaHeccoczS-",
        "outputId": "c13ce520-635c-4c97-ef0e-761a38147c8d"
      },
      "source": [
        "from jax import make_jaxpr\n",
        "\n",
        "trace_f = make_jaxpr(f)\n",
        "\n",
        "trace_f(3)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[]. let\n",
              "    b:i32[] = mul a 2\n",
              "    c:i32[] = mul b a\n",
              "    d:i32[] = mul a 3\n",
              "    e:i32[] = add c d\n",
              "    f:i32[] = add e 3\n",
              "  in (f,) }"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that this trace is done at a particular value. This is an important point; keep it at the back of your mind.\n",
        "\n",
        "Now this is just a function. But what is a function transformation? It is a function, that takes in a function, and pops out another function.\n",
        "\n",
        "Sunch transformations are at the core of functional programming. They provide a simple way to compose functions into a useful pieces of code.\n",
        "\n",
        "For example, consider:"
      ],
      "metadata": {
        "id": "XpBXhqWdrM5Z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def sum_of_anything(f):\n",
        "  def doit(a, b):\n",
        "    return f(a) + f(b)\n",
        "  return doit"
      ],
      "metadata": {
        "id": "Z_XKs6lZsL1n"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sum_of_squares = sum_of_anything(lambda x: x*x)\n",
        "sum_of_squares(3, 4)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PIjqZ4deslPF",
        "outputId": "b896f577-2686-4a9b-be21-bf0ca86fee45"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "25"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sum_of_cubes = sum_of_anything(lambda x: x*x*x)\n",
        "sum_of_cubes(3, 4)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_wxHnLxYsuxE",
        "outputId": "9aa9ec60-414f-4ea2-a4d2-4b4a2ff1d3c3"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "91"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "By returning a function we have created ageneral system that can be applied to any function of one variable, and return a sum.\n",
        "\n",
        "Jax implements four key function transformations for us which enable, in this functional style above, a lot of great functionality. These are (from https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/#references):\n",
        "\n",
        "1. `grad()` for evaluating the gradient function of the input function\n",
        "2. jit() to transform functions into just-in-time compiled versions\n",
        "3. vmap() for automatic vectorization of operations\n",
        "4. pmap() for easy parallelization of computations\n",
        "5. make_jaxpr() to get insight into what any jax function is doing\n",
        "\n",
        "Let's tackle these one by one and see how they compose for us a lovely deep learning and scientific computation system...\n"
      ],
      "metadata": {
        "id": "nohMhITCs8T1"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QMXGV3TjouWn"
      },
      "source": [
        "##  3: Automatic Differentiation and  Gradients\n",
        "\n",
        "There is a key difference between Jax and other deep learning frameworks like tensorflow and pytorch. Rather than compute the gradient at a loss function at a certain point by backpropogation through a computational graph, Jax implements a function transformation and returns the gradient function.\n",
        "\n",
        "What is a computational graph? Consider the trace we saw above:\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "trace_f(3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ajguY98Tu0Jn",
        "outputId": "60aa7ed4-021c-41ec-feb0-4592bfb90482"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[]. let\n",
              "    b:i32[] = mul a 2\n",
              "    c:i32[] = mul b a\n",
              "    d:i32[] = mul a 3\n",
              "    e:i32[] = add c d\n",
              "    f:i32[] = add e 3\n",
              "  in (f,) }"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This computation can be described by the following graph:\n",
        "\n",
        "![](https://i.imgur.com/V0y5oqF.jpg)\n",
        "\n",
        "Now, basic rules of calculus can be used to differentiate this graph. The question is, how is this done.\n",
        "\n",
        "`Pytorch` for example uses the original `autograd` library. As very nicely described in Sabrina Mielke's post, the rough idea there is:\n",
        "\n",
        "![](https://sjmielke.com/images/blog/jax-purify/comparison_small.png)\n",
        "\n",
        "As she describes:\n",
        "\n",
        ">PyTorch builds up a graph as you compute the forward pass, and one call to backward() on some “result” node then augments each intermediate node in the graph with the gradient of the result node with respect to that intermediate node. JAX on the other hand makes you express your computation as a Python function, and by transforming it with grad() gives you a gradient function that you can evaluate like your computation function—but instead of the output it gives you the gradient of the output with respect to (by default) the first parameter that your function took as input:"
      ],
      "metadata": {
        "id": "4P8JGQ_6u6WX"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eZ5XWCPrcpzn",
        "outputId": "ab1d49c0-a40f-48b5-a247-0b61448994ca"
      },
      "source": [
        "from jax import grad\n",
        "\n",
        "g_f = grad(f)\n",
        "print(type(g_f))\n",
        "g_f(3.)"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<class 'function'>\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(15., dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lfKYnsOSbddE"
      },
      "source": [
        "Currently (from a assumbly.ai aticle):\n",
        "\n",
        ">With grad(), you can differentiate through native Python and NumPy functions, such as loops, branches, recursion, closures, and “PyTrees” (e.g. dictionaries).\n",
        "\n",
        "`make_jaxpr` tells us what is happening in `grad`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QKT078mTdnqa",
        "outputId": "5741d777-5a9c-4209-f9aa-af3904aa4f6b"
      },
      "source": [
        "trace = make_jaxpr(grad(f))\n",
        "trace(3.)"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[]. let\n",
              "    b:f32[] = mul a 2.0\n",
              "    c:f32[] = mul b a\n",
              "    d:f32[] = mul a 3.0\n",
              "    e:f32[] = add c d\n",
              "    _:f32[] = add e 3.0\n",
              "    f:f32[] = mul 1.0 3.0\n",
              "    g:f32[] = mul b 1.0\n",
              "    h:f32[] = mul 1.0 a\n",
              "    i:f32[] = add_any f g\n",
              "    j:f32[] = mul h 2.0\n",
              "    k:f32[] = add_any i j\n",
              "  in (k,) }"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v7j4QV-WMb8K"
      },
      "source": [
        "`value_and_grad` gibe us the function as well asthe gradient:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I-P-ncDlMXFo",
        "outputId": "10d97135-2699-4409-a8b5-8619b9ba12ef"
      },
      "source": [
        "from jax import value_and_grad\n",
        "\n",
        "vg_f = value_and_grad(f)\n",
        "\n",
        "value, gradient = vg_f(3.)\n",
        "value, gradient"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(DeviceArray(30., dtype=float32, weak_type=True),\n",
              " DeviceArray(15., dtype=float32, weak_type=True))"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rM0ESzK2qSgI"
      },
      "source": [
        "## 4: `jit`ting\n",
        "\n",
        "We have seen `jit`ting before, but the summary is that [xla](https://www.tensorflow.org/xla) transforms functions from the python version to another one  with the same signature but which is optimised for the accelerator you are running on (e.g. GPU or TPU).\n",
        "\n",
        "The code is compiled at runtime, so the first execution will be slower.\n",
        "\n",
        "One of the key things that XLA will do for us is to fuse kernels: in other words, make code transformations that improve efficiency. This is done at the intermediate-representation (IR) level."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M9CUDRzosLqk",
        "outputId": "a20c4ac5-48c1-4133-b724-f8349b93309f"
      },
      "source": [
        "from jax import jit\n",
        "\n",
        "jitted_f = jit(f)\n",
        "\n",
        "jitted_f(3)"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(30, dtype=int32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "myU3A2HqFUYJ",
        "outputId": "9d2c4982-c506-49f7-9527-2f01d911fbb3"
      },
      "source": [
        "make_jaxpr(jitted_f)(3)"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[]. let\n",
              "    b:i32[] = xla_call[\n",
              "      call_jaxpr={ lambda ; c:i32[]. let\n",
              "          d:i32[] = mul c 2\n",
              "          e:i32[] = mul d c\n",
              "          f:i32[] = mul c 3\n",
              "          g:i32[] = add e f\n",
              "          h:i32[] = add g 3\n",
              "        in (h,) }\n",
              "      name=f\n",
              "    ] a\n",
              "  in (b,) }"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def ff(x,y):\n",
        "  return x**2 + y\n",
        "\n",
        "# Partial derviatives\n",
        "x = 2.0; y= 3.0;\n",
        "v, gx = value_and_grad(ff, argnums=0)(x,y)\n",
        "print(v)\n",
        "print(gx)\n",
        "\n",
        "gy = grad(ff, argnums=1)(x,y)\n",
        "print(gy)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VZEqsPtW7Q4z",
        "outputId": "e20925da-1367-459c-c34e-f262973a60cd"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "7.0\n",
            "4.0\n",
            "1.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can differentiate to any level. Higher order optimization routines ususally want hessians, which are very easy to calculate (more later)\n",
        "\n",
        "One thing you have probably been wondering about (because we alluded to it) is why these traces always have an argument.\n",
        "\n",
        "The answer lies in how `jit`ting is carried out. Remember that we first convert a function to the IR called jaxpr, and then the XLA JIT compiler will particularize the code for CPU, GPU, or TPU.\n",
        "\n",
        "This jaxpr is created by tracing the function for a specific value. So functions with wildly different behavior for different types, or conditional branches, need to be  treated more carefully. I'll defer to the documentation and https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/jax_intro.ipynb#scrollTo=6Ps1W8LhKKj9 for more details.\n",
        "\n"
      ],
      "metadata": {
        "id": "345MPy426_Mt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def f3(x):\n",
        "  if x > 0:\n",
        "    return x\n",
        "  else:\n",
        "    return 2 * x\n",
        "\n",
        "print(f(3.0))\n",
        "\n",
        "f3_jit = jit(f3)\n",
        "print(f3_jit(3.0))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 852
        },
        "id": "EEQWkuyA8kD3",
        "outputId": "a766ffc9-64c8-4260-e330-1c2d9eb08556"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "30.0\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "ConcretizationTypeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mUnfilteredStackTrace\u001b[0m                      Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-29-efe830384b61>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mf3_jit\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf3_jit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py\u001b[0m in \u001b[0;36mreraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    161\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 162\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    163\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mcache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    467\u001b[0m         \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbackend\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mflat_fun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 468\u001b[0;31m         donated_invars=donated_invars, inline=inline)\n\u001b[0m\u001b[1;32m    469\u001b[0m     \u001b[0mout_pytree_def\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout_tree\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1795\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1796\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mcall_bind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1797\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1811\u001b[0m   \u001b[0mfun_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mannotate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1812\u001b[0;31m   \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtop_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1813\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mapply_todos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_trace_todo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    680\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 681\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    682\u001b[0m   \u001b[0mprocess_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36m_xla_call_impl\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    149\u001b[0m   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001b[0;32m--> 150\u001b[0;31m                                *arg_specs)\n\u001b[0m\u001b[1;32m    151\u001b[0m   \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mmemoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    284\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 285\u001b[0;31m       \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    286\u001b[0m       \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mans\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36m_xla_callable_uncached\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m    197\u001b[0m   return lower_xla_callable(fun, device, backend, name, donated_invars, False,\n\u001b[0;32m--> 198\u001b[0;31m                             *arg_specs).compile().unsafe_call\n\u001b[0m\u001b[1;32m    199\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mTraceAnnotation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdecorator_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36mlower_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, always_lower, *arg_specs)\u001b[0m\n\u001b[1;32m    228\u001b[0m     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(\n\u001b[0;32m--> 229\u001b[0;31m         fun, abstract_args, pe.debug_info_final(fun, \"jit\"), which_explicit)\n\u001b[0m\u001b[1;32m    230\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTracer\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mTraceAnnotation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdecorator_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py\u001b[0m in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, debug_info, keep_inputs)\u001b[0m\n\u001b[1;32m   1854\u001b[0m       jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(\n\u001b[0;32m-> 1855\u001b[0;31m         fun, main, in_avals, keep_inputs=keep_inputs)\n\u001b[0m\u001b[1;32m   1856\u001b[0m     \u001b[0;32mdel\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py\u001b[0m in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals, keep_inputs)\u001b[0m\n\u001b[1;32m   1825\u001b[0m     \u001b[0min_tracers_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mt\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0min_tracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeep_inputs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mkeep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1826\u001b[0;31m     \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0min_tracers_\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1827\u001b[0m     \u001b[0mout_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mans\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    167\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 168\u001b[0;31m       \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    169\u001b[0m     \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-29-efe830384b61>\u001b[0m in \u001b[0;36mf3\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mf3\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0;32mif\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36m__bool__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    602\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m__nonzero__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_nonzero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 603\u001b[0;31m   \u001b[0;32mdef\u001b[0m \u001b[0m__bool__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_bool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    604\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m__int__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_int\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36merror\u001b[0;34m(self, arg)\u001b[0m\n\u001b[1;32m   1143\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1144\u001b[0;31m     \u001b[0;32mraise\u001b[0m \u001b[0mConcretizationTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfname_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1145\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\nThe problem arose with the `bool` function. \nWhile tracing the function f3 at <ipython-input-29-efe830384b61>:1 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.\n\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mConcretizationTypeError\u001b[0m                   Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-29-efe830384b61>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mf3_jit\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf3_jit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-29-efe830384b61>\u001b[0m in \u001b[0;36mf3\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mf3\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0;32mif\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0;36m2\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mConcretizationTypeError\u001b[0m: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\nThe problem arose with the `bool` function. \nWhile tracing the function f3 at <ipython-input-29-efe830384b61>:1 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.\n\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PGny96UJoqFp"
      },
      "source": [
        "## 5: Vectorization with `vmap`\n",
        "\n",
        "Jax provides you two places where vectorization is useful: takinf a function that acts on one input and making it act on more than one, and secondly, having the vectorization apply on a batch axis, so that the same function can be executed across multiple samples in a batch.\n",
        "\n",
        "Here is an example of the first kind of vectorization.\n",
        "We want to carry out a logistic regression prediction in this example from https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/jax_intro.ipynb#scrollTo=D2dSgG-cnVIe. Let us write the code to do this on a single example:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "D = 2\n",
        "N = 3\n",
        "\n",
        "w = np.random.normal(size=(D,))\n",
        "X = np.random.normal(size=(N,D))\n",
        "\n",
        "def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)\n",
        "\n",
        "def predict_single(x):\n",
        "    return sigmoid(jnp.dot(w, x)) # <(D) , (D)> = (1) # inner product"
      ],
      "metadata": {
        "id": "2acP6uonm9gV"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(predict_single(X[0,:])) # works"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yJ6BRZ0Tm__-",
        "outputId": "e2e5ca62-671f-4159-81f0-712f939ea142"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.37815177\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Clearly we cant use this code to predicr on a vector..."
      ],
      "metadata": {
        "id": "1tI_FbvEuMgv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(predict_single(X)) "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 520
        },
        "id": "6ZiQDyp5nD8D",
        "outputId": "4265d966-ae8e-48a7-87cc-4bdd7937e0d2"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "error",
          "ename": "TypeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mUnfilteredStackTrace\u001b[0m                      Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-21-3b4a74243467>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpredict_single\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-19-4bf036523ecc>\u001b[0m in \u001b[0;36mpredict_single\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mpredict_single\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0msigmoid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mw\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# <(D) , (D)> = (1) # inner product\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py\u001b[0m in \u001b[0;36mreraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    161\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 162\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    163\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mcache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    467\u001b[0m         \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mbackend\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mflat_fun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 468\u001b[0;31m         donated_invars=donated_invars, inline=inline)\n\u001b[0m\u001b[1;32m    469\u001b[0m     \u001b[0mout_pytree_def\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout_tree\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1795\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1796\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mcall_bind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1797\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1811\u001b[0m   \u001b[0mfun_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mannotate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1812\u001b[0;31m   \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtop_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1813\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mapply_todos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_trace_todo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    680\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 681\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    682\u001b[0m   \u001b[0mprocess_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36m_xla_call_impl\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    149\u001b[0m   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001b[0;32m--> 150\u001b[0;31m                                *arg_specs)\n\u001b[0m\u001b[1;32m    151\u001b[0m   \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mmemoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    284\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 285\u001b[0;31m       \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    286\u001b[0m       \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mans\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36m_xla_callable_uncached\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m    197\u001b[0m   return lower_xla_callable(fun, device, backend, name, donated_invars, False,\n\u001b[0;32m--> 198\u001b[0;31m                             *arg_specs).compile().unsafe_call\n\u001b[0m\u001b[1;32m    199\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mTraceAnnotation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdecorator_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py\u001b[0m in \u001b[0;36mlower_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, always_lower, *arg_specs)\u001b[0m\n\u001b[1;32m    228\u001b[0m     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(\n\u001b[0;32m--> 229\u001b[0;31m         fun, abstract_args, pe.debug_info_final(fun, \"jit\"), which_explicit)\n\u001b[0m\u001b[1;32m    230\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTracer\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mTraceAnnotation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdecorator_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py\u001b[0m in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, debug_info, keep_inputs)\u001b[0m\n\u001b[1;32m   1854\u001b[0m       jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(\n\u001b[0;32m-> 1855\u001b[0;31m         fun, main, in_avals, keep_inputs=keep_inputs)\n\u001b[0m\u001b[1;32m   1856\u001b[0m     \u001b[0;32mdel\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py\u001b[0m in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals, keep_inputs)\u001b[0m\n\u001b[1;32m   1825\u001b[0m     \u001b[0min_tracers_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mt\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0min_tracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeep_inputs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mkeep\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1826\u001b[0;31m     \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0min_tracers_\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1827\u001b[0m     \u001b[0mout_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mans\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    167\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 168\u001b[0;31m       \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    169\u001b[0m     \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36mdot\u001b[0;34m(a, b, precision)\u001b[0m\n\u001b[1;32m   2699\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0m_max\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma_ndim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_ndim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2700\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprecision\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mprecision\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2701\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/lax/lax.py\u001b[0m in \u001b[0;36mdot\u001b[0;34m(lhs, rhs, precision, preferred_element_type)\u001b[0m\n\u001b[1;32m    656\u001b[0m     raise TypeError(\"Incompatible shapes for dot: got {} and {}.\".format(\n\u001b[0;32m--> 657\u001b[0;31m         lhs.shape, rhs.shape))\n\u001b[0m\u001b[1;32m    658\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mUnfilteredStackTrace\u001b[0m: TypeError: Incompatible shapes for dot: got (2,) and (3, 2).\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-21-3b4a74243467>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpredict_single\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-19-4bf036523ecc>\u001b[0m in \u001b[0;36mpredict_single\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mpredict_single\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0msigmoid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mw\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# <(D) , (D)> = (1) # inner product\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36mdot\u001b[0;34m(a, b, precision)\u001b[0m\n\u001b[1;32m   2698\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2699\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0m_max\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma_ndim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_ndim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2700\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprecision\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mprecision\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2701\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2702\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mb_ndim\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mTypeError\u001b[0m: Incompatible shapes for dot: got (2,) and (3, 2)."
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "So we could do the matrix multiplication ourselves, taking care to match dimensions..."
      ],
      "metadata": {
        "id": "C5Oo4lCQuUqL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def predict_batch(X):\n",
        "    return sigmoid(jnp.dot(X, w)) # (N,D) * (D,1) = (N,1) # matrix-vector multiply\n",
        "\n",
        "print(predict_batch(X)) "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "33P9Eua5nI9S",
        "outputId": "3c182db8-c452-4d14-c855-c628e128dbc0"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.37815177 0.55771613 0.4303379 ]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Or let `jax` do it using `vmap`."
      ],
      "metadata": {
        "id": "l13sCHZuueUk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import vmap\n",
        "print(vmap(predict_single)(X))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Od5DwG2xnS4D",
        "outputId": "a0384eb8-1eb1-485a-de9e-1547e78207b5"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.37815177 0.55771613 0.4303379 ]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "`vmap` vectorizes over the first axis of each of its inputs. In the above example there is only one input, which is 3 samples of size 2. The first axis of this 2-D arrray is the 0th axis, or the sample axis. Which is exactly what we want in a prediction system.\n",
        "\n",
        "This way of thinking of the 0th axis as the sample axis, or the batch axis leads to the second usage of `vmap`. Which is to carry out ops on everything but the sample/batch axis. Consider the example below (from http://matpalm.com/blog/ymxb_pod_slice/):"
      ],
      "metadata": {
        "id": "9rxzs6_xujhy"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tPHAWQKeepDV"
      },
      "source": [
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "\n",
        "def f(x):    \n",
        "  return jnp.array([jnp.min(x), jnp.max(x)])"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "diIx82dwvLhw"
      },
      "source": [
        "Suppose our input has shape: `(2, 3)`. Then we'll get the min and the max in the whole array:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "o38zJOxghIJE",
        "outputId": "2d9e6292-d13b-4d79-c603-a33d7ae72b92"
      },
      "source": [
        "x = np.arange(6).reshape((2, 3))\n",
        "x"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0, 1, 2],\n",
              "       [3, 4, 5]])"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V1D95TfWhJ7L",
        "outputId": "4ee3e055-8953-492f-dd27-49017c9ebc6d"
      },
      "source": [
        "f(x)"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([0, 5], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XeyzO5tZhU7M"
      },
      "source": [
        "Now, what if we have a batch of 4 inputs; i.e. x is of shape `(4, 2, 3)` and thus we want to return shape `(4, 2)`. Let us set up some data for this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wp-du_O5eq5c",
        "outputId": "4b548f9f-e4e8-4437-b2df-f61f90ef7438"
      },
      "source": [
        "bx = np.arange(24).reshape((4, 2, 3))\n",
        "bx"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[[ 0,  1,  2],\n",
              "        [ 3,  4,  5]],\n",
              "\n",
              "       [[ 6,  7,  8],\n",
              "        [ 9, 10, 11]],\n",
              "\n",
              "       [[12, 13, 14],\n",
              "        [15, 16, 17]],\n",
              "\n",
              "       [[18, 19, 20],\n",
              "        [21, 22, 23]]])"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Sixo-2chaYT"
      },
      "source": [
        "If you now call `f(x)` its gonna give you the min and max of the whole array, and thats not what you want to do. Normally you might use an explicit axis argument in the `min` and `max`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zn-qsoHnfXHF",
        "outputId": "e9d74c74-0aca-4da3-ce6d-a4fa0afd8e11"
      },
      "source": [
        "f(bx)"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([ 0, 23], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6akvgnThuh0"
      },
      "source": [
        "We can do it implicitly using `vmap`. The function you feed to `vmap` is expected to have arguments operating with a extra leading gimension. Normally this is the first argument..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rFYcb6rle1q_",
        "outputId": "1c819f6d-f4e2-488e-b3ff-db4081751619"
      },
      "source": [
        "from jax import vmap\n",
        "\n",
        "               # f  in:    (2, 3) out:    (2)\n",
        "v_f = vmap(f)  # vf in: (B, 2, 3) out: (B, 2)\n",
        "v_f(bx)"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[ 0,  5],\n",
              "             [ 6, 11],\n",
              "             [12, 17],\n",
              "             [18, 23]], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UZpizPXhjN9j",
        "outputId": "c1fce2ab-22ca-4400-9d36-6adb934c9ccc"
      },
      "source": [
        "make_jaxpr(f)(x)"
      ],
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[2,3]. let\n",
              "    b:i32[] = reduce_min[axes=(0, 1)] a\n",
              "    c:i32[] = reduce_max[axes=(0, 1)] a\n",
              "    d:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] b\n",
              "    e:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] c\n",
              "    f:i32[2] = concatenate[dimension=0] d e\n",
              "  in (f,) }"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E2qarXdDjW0x",
        "outputId": "771230db-daf5-48f6-f92c-ae2eb5d7420c"
      },
      "source": [
        "make_jaxpr(v_f)(bx)"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[4,2,3]. let\n",
              "    b:i32[4] = reduce_min[axes=(1, 2)] a\n",
              "    c:i32[4] = reduce_max[axes=(1, 2)] a\n",
              "    d:i32[4,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(4, 1)] b\n",
              "    e:i32[4,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(4, 1)] c\n",
              "    f:i32[4,2] = concatenate[dimension=1] d e\n",
              "  in (f,) }"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can now `jit` the function if you like to speed it up."
      ],
      "metadata": {
        "id": "vgmqx-UNAIFx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jv_f = jit(vmap(f))\n",
        "jv_f(bx)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y4FQoGIZaak3",
        "outputId": "21b3da58-8ae3-4210-cb87-0038762b819f"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[ 0,  5],\n",
              "             [ 6, 11],\n",
              "             [12, 17],\n",
              "             [18, 23]], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "make_jaxpr(jv_f)(bx)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IR48g0hVasSt",
        "outputId": "d055db6b-0d8e-43f5-8b16-255a7eabb05e"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[4,2,3]. let\n",
              "    b:i32[4,2] = xla_call[\n",
              "      call_jaxpr={ lambda ; c:i32[4,2,3]. let\n",
              "          d:i32[4] = reduce_min[axes=(1, 2)] c\n",
              "          e:i32[4] = reduce_max[axes=(1, 2)] c\n",
              "          f:i32[4,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(4, 1)] d\n",
              "          g:i32[4,1] = broadcast_in_dim[broadcast_dimensions=(0,) shape=(4, 1)] e\n",
              "          h:i32[4,2] = concatenate[dimension=1] f g\n",
              "        in (h,) }\n",
              "      name=f\n",
              "    ] a\n",
              "  in (b,) }"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIXHZ4msvs4A"
      },
      "source": [
        "As `f` and your tensors get more complex it is harder to make `f` batch-aware. `vmap` does it for you.\n",
        "\n",
        "The default behavior of `vmap` is to vectorise  across the first axis of the first arg, but we have a lot of control on how we want to vectorise; e.g. consider"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def f(a, b, c):\n",
        "  return a + b + c\n",
        "a = np.arange(6).reshape((2, 3)) + 10\n",
        "b = 2\n",
        "c = np.arange(6).reshape((3, 2)) + 20\n",
        "print(repr(a))\n",
        "c"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jirzYWszBKcR",
        "outputId": "9703604e-842a-4b45-8c8b-e1d4d5a2b864"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "array([[10, 11, 12],\n",
            "       [13, 14, 15]])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[20, 21],\n",
              "       [22, 23],\n",
              "       [24, 25]])"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v_f = vmap(f, in_axes=(0, None, 1))\n",
        "\n",
        "v_f(a, b, c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-nF6pRk4G7ZW",
        "outputId": "816540b0-120c-4e91-c0fb-7455df4138e6"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[32, 35, 38],\n",
              "             [36, 39, 42]], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ok, so what happened here? Look at the incantation of `in_axes`. The tuple there refers to the argument of `f`. `a` is a 2x3 so we want to vmap over the rows. with the tuple element being the axis 0. For `b` the tuple element is `None`, which means, dont vmap over this. Now for `c`, which is 3x2 it wants us to vmap over axis 1, or the horizontal axis. This means that the 22 in `c` is added to the 11 in `a` and then added to the broadcasted 2 to create the 35 in the output. You can see these machinations in the jaxpr for this function."
      ],
      "metadata": {
        "id": "OiC2m-YvHAFh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "make_jaxpr(v_f)(a, b, c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZDQfiDndJOAR",
        "outputId": "133a2776-aab7-436c-89f1-bbcc756ed1a6"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:i32[2,3] b:i32[] c:i32[3,2]. let\n",
              "    d:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b\n",
              "    e:i32[2,3] = add a d\n",
              "    f:i32[2,3] = transpose[permutation=(1, 0)] c\n",
              "    g:i32[2,3] = add e f\n",
              "  in (g,) }"
            ]
          },
          "metadata": {},
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HmBYDM_FnVC9"
      },
      "source": [
        "## 6: Parallelisation\n",
        "\n",
        "`pmap` has pretty much the same interface as  `vmap`. So why a new construct? Well, pmap uses XLA and provides automatic jitting. It does not provide a vectorized version like `vmap` does though. Instead, XLA is used to create a SPMD (Single Program Multiple Data) program to run the function in parallel across all available devices: CPU's, GPU's or TPUs.\n",
        "\n",
        "Here are devices are:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HvHcOpvQnirv",
        "outputId": "69d572fd-14a7-4ab5-e121-7ff626410d52"
      },
      "source": [
        "jax.devices()"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
              " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
              " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
              " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b8rz5cjQnoPU",
        "outputId": "9c4ba101-0a42-4a15-b14d-a83c2d764b3b"
      },
      "source": [
        "from jax import pmap\n",
        "\n",
        "def f(a, b, c):\n",
        "  return a + b + c\n",
        "  \n",
        "p_f = pmap(f, in_axes=(0, None, 0)) \n",
        "\n",
        "a = np.random.random(size=(num_replicas, 3))\n",
        "b = 4\n",
        "c = np.random.random(size=(num_replicas, 3))\n",
        "\n",
        "a"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.77718412, 0.00176833, 0.44951046],\n",
              "       [0.27297925, 0.46497406, 0.91643829],\n",
              "       [0.09127398, 0.59443344, 0.95430272],\n",
              "       [0.00510437, 0.10622125, 0.29679499],\n",
              "       [0.78516484, 0.09587237, 0.7798978 ],\n",
              "       [0.50380788, 0.25051366, 0.61133905],\n",
              "       [0.56083878, 0.44882798, 0.2974137 ],\n",
              "       [0.96834715, 0.72368435, 0.94016945]])"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice now that we have as many rows as we have devices. Now it is not the batch axis we are \"vectorizing\" or more precisely \"spmd\"ing over, but rather this device dimension"
      ],
      "metadata": {
        "id": "rmlDIEG2M92H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "p_f(a, b, c)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pGvn-fDcNQ9i",
        "outputId": "60ad5ed6-7ef8-4d32-d30a-479d71a90c43"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "ShardedDeviceArray([[5.010401 , 4.144878 , 4.6788063],\n",
              "                    [5.11807  , 4.5080056, 5.2585855],\n",
              "                    [4.3505087, 4.9883356, 5.7644   ],\n",
              "                    [4.076356 , 4.7490406, 4.4232683],\n",
              "                    [4.9303174, 4.374606 , 5.433234 ],\n",
              "                    [5.1140704, 4.5758166, 5.3989067],\n",
              "                    [5.4698825, 5.1499176, 4.4592896],\n",
              "                    [5.0957966, 5.681382 , 5.3702273]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that an identical program has been shipped off to 8 devices. A broadcasted addition has been carried out on these `num_replicas` devices. Notice the output data type. It is a _Sharded_ device array. This means that those parts of the array are not actually on the main machine, but rather on the individual devices. We have moved the program to the datas, and run it on each of the datas.\n",
        "\n",
        "jaxpr as usual tells us more:"
      ],
      "metadata": {
        "id": "scegIipQNY9o"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mxIN0WxUStKX",
        "outputId": "f2e9f963-2d6e-4519-be50-9a9acc4bd1a9"
      },
      "source": [
        "make_jaxpr(pmap(f, in_axes=(0, None, 0)))(a, b, c)"
      ],
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[8,3] b:i32[] c:f32[8,3]. let\n",
              "    d:f32[8,3] = xla_pmap[\n",
              "      axis_name=<axis 0x7f92facb9ef0>\n",
              "      axis_size=8\n",
              "      backend=None\n",
              "      call_jaxpr={ lambda ; e:f32[3] f:i32[] g:f32[3]. let\n",
              "          h:f32[] = convert_element_type[new_dtype=float32 weak_type=False] f\n",
              "          i:f32[3] = add e h\n",
              "          j:f32[3] = add i g\n",
              "        in (j,) }\n",
              "      devices=None\n",
              "      donated_invars=(False, False, False)\n",
              "      global_arg_shapes=(None, None, None)\n",
              "      global_axis_size=None\n",
              "      in_axes=(0, None, 0)\n",
              "      name=f\n",
              "      out_axes=(0,)\n",
              "    ] a b c\n",
              "  in (d,) }"
            ]
          },
          "metadata": {},
          "execution_count": 45
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iKOf8WE1oHHe"
      },
      "source": [
        "Our output is stored `ShardedDeviceArray`, which is an abstraction over remote arrays constructed on each device. Imagine a larger calculation where we want both input and output to be sharded across devices so there is minimum data transfer. This is classic compute at data SPMD."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GqJ07r06oewh"
      },
      "source": [
        "## 7: Combing `vmap` and `pmap`: composition again\n",
        "\n",
        "Here is a slightly different f: we matrix multiply and add:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zYZ5A9VGoeQ8",
        "outputId": "89009426-c7d0-4bbe-8caa-9a530a5e70be"
      },
      "source": [
        "def f(a, b, c):\n",
        "  return a @ b + c\n",
        "\n",
        "a = np.random.random(size=(2, 3))\n",
        "b = np.random.random(size=(3, 4))\n",
        "c = np.random.random(size=(2, 4))\n",
        "\n",
        "f(a, b, c)"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.76190696, 1.01631292, 0.83752259, 1.87883684],\n",
              "       [1.28544108, 1.36528368, 0.84584105, 1.27333203]])"
            ]
          },
          "metadata": {},
          "execution_count": 48
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yiOJUOCJ7dZR"
      },
      "source": [
        "Let us add a batch dimension and vmap."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "s5QGNE4AokoF",
        "outputId": "077b2f8e-b670-47c8-d66b-859a035d8276"
      },
      "source": [
        "v_f = vmap(f)\n",
        "\n",
        "a = np.random.random(size=(5, 2, 3))\n",
        "b = np.random.random(size=(5, 3, 4))\n",
        "c = np.random.random(size=(5, 2, 4))\n",
        "\n",
        "v_f(a, b, c)"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[[1.0452026 , 1.3957614 , 0.8974083 , 0.2570575 ],\n",
              "              [1.3226345 , 1.6668731 , 0.8013219 , 1.0267715 ]],\n",
              "\n",
              "             [[0.42478555, 0.48035237, 1.0448409 , 1.054654  ],\n",
              "              [1.2317226 , 1.3148572 , 0.54127437, 0.699142  ]],\n",
              "\n",
              "             [[1.8477    , 1.4496236 , 1.5217985 , 2.1509185 ],\n",
              "              [0.96027756, 0.20201734, 1.2771938 , 1.2706423 ]],\n",
              "\n",
              "             [[1.3032501 , 1.8480506 , 1.5902729 , 0.78827995],\n",
              "              [0.9036236 , 0.8725722 , 1.3766314 , 0.88759273]],\n",
              "\n",
              "             [[0.6759734 , 1.483907  , 0.9024238 , 1.6253147 ],\n",
              "              [0.5552775 , 2.0771217 , 0.7588869 , 1.6169556 ]]],            dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S73-yW0n72oS"
      },
      "source": [
        "Now we have 5 2x4 arrays. Lets jit to make sure we can run on an accelerator..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K1QTlSXc79rI",
        "outputId": "2a731ed3-a3a6-4a00-f45c-2cfdd08d36b7"
      },
      "source": [
        "jv_f = jit(vmap(f))\n",
        "\n",
        "a = np.random.random(size=(5, 2, 3))\n",
        "b = np.random.random(size=(5, 3, 4))\n",
        "c = np.random.random(size=(5, 2, 4))\n",
        "\n",
        "jv_f(a, b, c).shape"
      ],
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, 2, 4)"
            ]
          },
          "metadata": {},
          "execution_count": 50
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mxLuNyVw7rb7"
      },
      "source": [
        "What happens is we wrap the `vmap` with a `pmap`? We have all these tpus or cpus we want to use..\n",
        "\n",
        "Lets add a batch dimension which makes it useful to do something like this..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7SoNK-1Ho6Ei",
        "outputId": "be256bee-54ea-4de0-ecb8-16d6d6434879"
      },
      "source": [
        "pjv_f = pmap(vmap(f))\n",
        "\n",
        "a = np.random.random(size=(num_replicas, 5, 2, 3))\n",
        "b = np.random.random(size=(num_replicas, 5, 3, 4))\n",
        "c = np.random.random(size=(num_replicas, 5, 2, 4))\n",
        "\n",
        "pjv_f(a, b, c)"
      ],
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "ShardedDeviceArray([[[[2.1330729 , 1.079713  , 1.6747804 , 2.4263163 ],\n",
              "                      [0.72165257, 0.5211822 , 0.4666624 , 1.4015088 ]],\n",
              "\n",
              "                     [[1.4648497 , 1.58633   , 1.068167  , 1.4730864 ],\n",
              "                      [1.5538089 , 1.2625966 , 1.2007633 , 1.5472561 ]],\n",
              "\n",
              "                     [[1.3659166 , 0.9606316 , 1.2130721 , 0.9296636 ],\n",
              "                      [1.4485664 , 1.824628  , 0.9115788 , 0.98683596]],\n",
              "\n",
              "                     [[1.2636752 , 0.94934475, 0.9281723 , 1.1574597 ],\n",
              "                      [1.4386181 , 1.0023386 , 0.80278707, 0.8467031 ]],\n",
              "\n",
              "                     [[0.7390131 , 1.171852  , 0.98206335, 0.5656377 ],\n",
              "                      [0.91225004, 2.0936654 , 1.1719465 , 1.3942947 ]]],\n",
              "\n",
              "\n",
              "                    [[[0.7920339 , 1.0567851 , 1.3067839 , 0.8583878 ],\n",
              "                      [1.1420336 , 1.4154868 , 1.1066592 , 1.0678368 ]],\n",
              "\n",
              "                     [[0.49421442, 1.4615567 , 1.4088886 , 1.3548925 ],\n",
              "                      [0.5242963 , 1.0997252 , 1.3237301 , 1.4823478 ]],\n",
              "\n",
              "                     [[1.0898659 , 0.9365779 , 1.2225327 , 1.8666003 ],\n",
              "                      [2.1461802 , 1.2673559 , 2.1495433 , 2.4694262 ]],\n",
              "\n",
              "                     [[0.9406155 , 1.7621238 , 1.1258657 , 1.2252997 ],\n",
              "                      [0.69291174, 1.3035141 , 0.332399  , 1.1014758 ]],\n",
              "\n",
              "                     [[1.5623498 , 2.5511923 , 0.9915999 , 1.8267992 ],\n",
              "                      [0.6319851 , 1.7797315 , 1.6541996 , 1.5773919 ]]],\n",
              "\n",
              "\n",
              "                    [[[1.3744913 , 1.748306  , 0.8734106 , 0.6496623 ],\n",
              "                      [2.3969107 , 2.0945878 , 1.6375959 , 1.8564918 ]],\n",
              "\n",
              "                     [[1.4100482 , 1.6448351 , 1.012776  , 1.8215333 ],\n",
              "                      [1.1853483 , 1.6816251 , 1.1982111 , 1.0506394 ]],\n",
              "\n",
              "                     [[0.7229642 , 1.503804  , 1.5507644 , 0.7743245 ],\n",
              "                      [0.9667979 , 1.6087203 , 1.3655764 , 1.5802486 ]],\n",
              "\n",
              "                     [[1.1642957 , 1.5026107 , 1.0496005 , 1.1637139 ],\n",
              "                      [0.70106786, 1.2086004 , 1.2789868 , 1.2088027 ]],\n",
              "\n",
              "                     [[1.3797216 , 0.83027846, 1.4625288 , 2.1083858 ],\n",
              "                      [1.0929985 , 0.7565248 , 1.0053297 , 0.89984506]]],\n",
              "\n",
              "\n",
              "                    [[[0.6196115 , 1.3597128 , 1.3710406 , 1.5466938 ],\n",
              "                      [0.93426436, 0.8885168 , 2.180446  , 1.488902  ]],\n",
              "\n",
              "                     [[0.8220608 , 0.6567021 , 1.0518647 , 1.6454827 ],\n",
              "                      [1.1435785 , 0.77745277, 1.595123  , 1.2485783 ]],\n",
              "\n",
              "                     [[1.2667642 , 1.3670704 , 1.1986128 , 1.3832426 ],\n",
              "                      [0.99903464, 0.71794224, 1.0667812 , 1.346477  ]],\n",
              "\n",
              "                     [[0.9903511 , 0.9011689 , 1.1929986 , 1.1412956 ],\n",
              "                      [2.0482147 , 2.2433264 , 1.1581957 , 1.6074884 ]],\n",
              "\n",
              "                     [[1.597811  , 2.00544   , 1.4720745 , 1.744862  ],\n",
              "                      [1.552243  , 2.150023  , 1.258941  , 1.3934748 ]]],\n",
              "\n",
              "\n",
              "                    [[[0.8211899 , 1.2530806 , 0.5969001 , 0.954087  ],\n",
              "                      [0.43300724, 1.433332  , 1.1253145 , 1.476504  ]],\n",
              "\n",
              "                     [[1.5236706 , 1.5442433 , 1.4183376 , 2.1751158 ],\n",
              "                      [0.84046245, 0.53063947, 0.6858475 , 0.99108166]],\n",
              "\n",
              "                     [[1.0214491 , 1.5879257 , 0.83864754, 0.5717579 ],\n",
              "                      [0.83486736, 1.7954059 , 1.7656362 , 1.1552426 ]],\n",
              "\n",
              "                     [[1.2005183 , 0.5581429 , 1.2542965 , 1.8650296 ],\n",
              "                      [1.3108244 , 1.0237314 , 1.1274173 , 0.7676884 ]],\n",
              "\n",
              "                     [[1.3998419 , 0.6368445 , 0.2114273 , 0.73516977],\n",
              "                      [1.2262338 , 0.63622874, 0.55001724, 0.30595717]]],\n",
              "\n",
              "\n",
              "                    [[[1.1612862 , 0.90250176, 0.72555614, 1.1632867 ],\n",
              "                      [1.1244104 , 1.546227  , 1.6493639 , 2.2057867 ]],\n",
              "\n",
              "                     [[0.4515928 , 1.9257517 , 0.7144293 , 1.1708785 ],\n",
              "                      [1.2895532 , 1.604795  , 0.8666728 , 0.9277959 ]],\n",
              "\n",
              "                     [[1.3056563 , 1.0159234 , 1.8094538 , 1.4023085 ],\n",
              "                      [1.0397915 , 1.2523758 , 1.5802643 , 1.3335234 ]],\n",
              "\n",
              "                     [[0.84180105, 1.6095195 , 1.4934287 , 1.3692423 ],\n",
              "                      [0.22922206, 1.2786512 , 1.0255628 , 0.5058948 ]],\n",
              "\n",
              "                     [[1.843193  , 1.4342294 , 1.7119837 , 1.8501928 ],\n",
              "                      [1.7453145 , 0.7689618 , 1.5866386 , 1.9972321 ]]],\n",
              "\n",
              "\n",
              "                    [[[0.9577383 , 1.481104  , 1.7567213 , 1.0351063 ],\n",
              "                      [1.9336603 , 1.2933993 , 1.6251562 , 2.325394  ]],\n",
              "\n",
              "                     [[2.123846  , 2.46204   , 1.7202034 , 1.4234455 ],\n",
              "                      [1.7550598 , 1.2197073 , 1.7540308 , 1.3295074 ]],\n",
              "\n",
              "                     [[0.66353583, 1.0090213 , 1.2509401 , 0.9847655 ],\n",
              "                      [1.3182061 , 1.2576656 , 1.8444761 , 1.1926347 ]],\n",
              "\n",
              "                     [[1.623388  , 1.8680383 , 0.2751102 , 0.92558765],\n",
              "                      [1.9880342 , 1.1245443 , 1.3466474 , 1.3849825 ]],\n",
              "\n",
              "                     [[1.0889342 , 2.148755  , 1.4684262 , 1.6300986 ],\n",
              "                      [0.48554766, 1.7777543 , 1.6985788 , 1.1625669 ]]],\n",
              "\n",
              "\n",
              "                    [[[1.7262106 , 1.5194578 , 0.7072493 , 1.8636508 ],\n",
              "                      [0.92681956, 0.72600496, 0.57179004, 1.2066619 ]],\n",
              "\n",
              "                     [[1.7491584 , 1.167714  , 1.354742  , 0.8789219 ],\n",
              "                      [1.2225047 , 0.6069238 , 1.2081804 , 1.1919103 ]],\n",
              "\n",
              "                     [[1.4697517 , 1.2906985 , 0.9590897 , 1.9893222 ],\n",
              "                      [1.0349257 , 1.6041591 , 1.5248736 , 1.15854   ]],\n",
              "\n",
              "                     [[0.7408388 , 0.74209183, 1.0904496 , 1.0586886 ],\n",
              "                      [0.21633269, 0.9424185 , 0.6659249 , 1.2457364 ]],\n",
              "\n",
              "                     [[1.1803762 , 1.7383925 , 1.0098609 , 1.485389  ],\n",
              "                      [1.5196466 , 1.336801  , 0.6660111 , 0.36105847]]]],                   dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Each vmap gave us 5 arrays. But because we had `num_replicas` of these, each got assigned to a compute device, and the calculation got garried out there. In this way we can reap the advantages of both vectorization and  parallelism."
      ],
      "metadata": {
        "id": "FE0e6VCygdJc"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eu_tshW-aiQR"
      },
      "source": [
        "\n",
        "\n",
        "## 8. State: random numbers in jax\n",
        "\n",
        "Jax is a pure functional framework. It ont allow you to change `DeviceArrays`, as opposed to how numpy deals with `ndarrays`.\n",
        "\n",
        "This means that jax never mutates arrays in place. If you are familiar with Spark dataframes, you will find the same restriction.\n",
        "\n",
        "This restriction causes us to write functiona with _referential transparancy_. Any time you run these functions, you get the samew result. State cannot play a role.\n",
        "\n",
        "This is important as it is precisely what enables the spark schedulers or xla compilers wo wotl. When you know that a piece of data or metadata is not going to change except be explicit change rather than by some object oriented hidden state, you can compose tranformations and make other efficiency changes. As well, your computation can be statically, state-free analysed. It is a radeoff..you will create new arrays rather than mutate them in place, requiring more memory. But it also means you can reason about and transform your computation.\n",
        "\n",
        "One os the huge epitomes of state in a python program is the random number generator.\n",
        "\n",
        "Jax allows stateless random number generayion, but it thus requires more care than in the numpy case.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IymD8avuKuf_",
        "outputId": "7646593b-0d3a-4bfa-f445-9bee9f6f5a8a"
      },
      "source": [
        "key = jax.random.PRNGKey(1337)\n",
        "print(jax.random.uniform(key))\n",
        "print(jax.random.uniform(key))"
      ],
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.02251327\n",
            "0.02251327\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UWaH5XODK1SY"
      },
      "source": [
        "when you want to generate more random numbers you need to explicit split the key"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T1K7PJUuKteL",
        "outputId": "9909a34b-3620-4e47-91d8-44d00e30b084"
      },
      "source": [
        "key = jax.random.PRNGKey(1337)\n",
        "for _ in range(10):\n",
        "  key, key2 = jax.random.split(key)\n",
        "  print(jax.random.uniform(key2))"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.31418657\n",
            "0.8795924\n",
            "0.81679654\n",
            "0.82754505\n",
            "0.78958607\n",
            "0.9006636\n",
            "0.18430483\n",
            "0.2509787\n",
            "0.5041659\n",
            "0.11195123\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# You cannot assign directly to elements of an array.\n",
        "\n",
        "A = jnp.zeros((3,3), dtype=np.float32)\n",
        "\n",
        "# In place update of JAX's array will yield an error!\n",
        "try:\n",
        "  A[1, :] = 1.0\n",
        "except:\n",
        "  print('cannot update in place')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-3WeQYAqoYWb",
        "outputId": "e394a001-17ed-4c5c-aa87-05dc5d568def"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cannot update in place\n"
          ]
        }
      ]
    }
  ]
}