Chapter-7/JAX_in_Action_Chapter_7_pmap.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCirjpcFSPan"
      },
      "source": [
        "## ⚖️ Choose A or B:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5lNqNEGtLyO"
      },
      "source": [
        "## A: Emulating multi-device system on CPU\n",
        "\n",
        "Use this section to initialize a set of virtual devices on CPU if you have no access to a multi-device system.\n",
        "\n",
        "It can also help you prototype, debug and test your multi-device code locally before running it on the expensive system.\n",
        "\n",
        "Even in the case of using Google Colab it can help you prototype faster because a CPU runtime is faster to restart."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ypaH8OgftR_H"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n8ip1VKgterO"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tlh75iWutfVT",
        "outputId": "1e3f0edd-2820-445d-ce64-c095f889e2b7"
      },
      "outputs": [],
      "source": [
        "jax.devices(\"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rmk_qf6mKCZx"
      },
      "source": [
        "## B: Setting up TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RMq70OC9wVHQ"
      },
      "source": [
        "!! **First, use instruction from Appendix C or example from Chapter 3 to run Cloud TPU and connect Colab to it.** Then run code below !!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YnSeGp15SDKi"
      },
      "outputs": [],
      "source": [
        "# install if you didn't install it on a Cloud TPU VM\n",
        "!pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FDq6YUpdSLNf",
        "outputId": "59875e4e-4fd8-4db6-a238-32ff0f8d33b9"
      },
      "outputs": [],
      "source": [
        "# need for our later example with SPMD neural net training\n",
        "!pip install tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VMCFxcZcS0u0",
        "outputId": "d042dd40-41ec-4ccc-dfb9-dd913035d37a"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2aBy_QGmu-61",
        "outputId": "42573514-97a2-43b5-bd84-f29ac4a15cb8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tpu\n"
          ]
        }
      ],
      "source": [
        "from jax.lib import xla_bridge\n",
        "print(xla_bridge.get_backend().platform)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "7c2IXWjRvFbj"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_qLo4vNRvHc-",
        "outputId": "f1a64216-4f71-4dc0-b052-6405b63aab8a"
      },
      "outputs": [
        {
          "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=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wzZlvWsjKIC8"
      },
      "source": [
        "## Parallelizing a function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "qxtvsXER1Ums"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UuEi71y-1iDg",
        "outputId": "2d1a4f77-51da-4e3c-ea3a-b39c069ec2c8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(2., dtype=float32)"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dot(jnp.array([1., 1., 1.]), jnp.array([1., 2., -1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6uqvbVsKLdp"
      },
      "source": [
        "Generate some **large** random arrays"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 141,
      "metadata": {
        "id": "4cVOwuJyv6yO"
      },
      "outputs": [],
      "source": [
        "from jax import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 142,
      "metadata": {
        "id": "EUTNwutCwiUl"
      },
      "outputs": [],
      "source": [
        "rng_key = random.PRNGKey(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 143,
      "metadata": {
        "id": "MSGxkzRGwmiy"
      },
      "outputs": [],
      "source": [
        "vs = random.normal(rng_key, shape=(20_000_000,3))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 144,
      "metadata": {
        "id": "zLYOfhF19gI4"
      },
      "outputs": [],
      "source": [
        "v1s = vs[:10_000_000,:]\n",
        "v2s = vs[10_000_000:,:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 145,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oBMFv_cc-GV6",
        "outputId": "a6530e75-e9e9-44b8-cebb-f0a0b82ed8db"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((10000000, 3), (10000000, 3))"
            ]
          },
          "execution_count": 145,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fHT8Pi7YPNvC"
      },
      "source": [
        "Make a compiled vectorized function as a baseline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 146,
      "metadata": {
        "id": "KsRK7fsv92RN"
      },
      "outputs": [],
      "source": [
        "dot_batched = jax.jit(jax.vmap(dot)).lower(v1s,v2s).compile() # you can use AOT-compilation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 147,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b-NUKyM6gO1R",
        "outputId": "a82076bc-4a26-4e92-c827-7cd762f5eda7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "503 μs ± 994 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit x_vmap = dot_batched(v1s, v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 148,
      "metadata": {
        "id": "yPKa12rgEan4"
      },
      "outputs": [],
      "source": [
        "dot_batched = jax.jit(jax.vmap(dot)) # or JIT-compilation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 149,
      "metadata": {
        "id": "tQd2syYJllpQ"
      },
      "outputs": [],
      "source": [
        "x_vmap = dot_batched(v1s, v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 150,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9uZUaB6nE5pT",
        "outputId": "6a81e3a1-ca18-4543-e786-f6eb13a5d1cb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "execution_count": 150,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_vmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 151,
      "metadata": {
        "id": "ZkmFPNA-hjbR"
      },
      "outputs": [],
      "source": [
        "dot_parallel = jax.pmap(dot)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 152,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 488
        },
        "id": "J4s0rZIVh0RC",
        "outputId": "6301885b-5859-45ac-ac53-d76a18ecb813"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "compiling computation that requires 10000000 logical devices, but only 4 XLA devices are available (num_replicas=10000000)",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[152], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m x_pmap \u001b[38;5;241m=\u001b[39m \u001b[43mdot_parallel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1s\u001b[49m\u001b[43m,\u001b[49m\u001b[43mv2s\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
            "File \u001b[0;32m~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:974\u001b[0m, in \u001b[0;36mUnloadedPmapExecutable.from_hlo\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    971\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m shards\u001b[38;5;241m.\u001b[39mnum_global_shards \u001b[38;5;241m>\u001b[39m xb\u001b[38;5;241m.\u001b[39mdevice_count(pci\u001b[38;5;241m.\u001b[39mbackend):\n\u001b[1;32m    972\u001b[0m   msg \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompiling computation that requires \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m logical devices, but only \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m XLA \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    973\u001b[0m          \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdevices are available (num_replicas=\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 974\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg\u001b[38;5;241m.\u001b[39mformat(shards\u001b[38;5;241m.\u001b[39mnum_global_shards,\n\u001b[1;32m    975\u001b[0m                               xb\u001b[38;5;241m.\u001b[39mdevice_count(pci\u001b[38;5;241m.\u001b[39mbackend),\n\u001b[1;32m    976\u001b[0m                               replicas\u001b[38;5;241m.\u001b[39mnum_global_replicas))\n\u001b[1;32m    977\u001b[0m \u001b[38;5;66;03m# On a single host, we simply grab the first N devices from jax.devices().\u001b[39;00m\n\u001b[1;32m    978\u001b[0m \u001b[38;5;66;03m# In the single host case, we want the default device order of pmap to\u001b[39;00m\n\u001b[1;32m    979\u001b[0m \u001b[38;5;66;03m# match jax.devices().\u001b[39;00m\n\u001b[1;32m    980\u001b[0m \u001b[38;5;66;03m# On multiple hosts, we create a default device assignment that ensures\u001b[39;00m\n\u001b[1;32m    981\u001b[0m \u001b[38;5;66;03m# each host is responsible for a contiguous set of replicas.\u001b[39;00m\n\u001b[1;32m    982\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m shards\u001b[38;5;241m.\u001b[39mnum_global_shards \u001b[38;5;241m>\u001b[39m shards\u001b[38;5;241m.\u001b[39mnum_local_shards:\n\u001b[1;32m    983\u001b[0m   \u001b[38;5;66;03m# TODO(skye): use a locality-aware assignment that satisfies the above\u001b[39;00m\n\u001b[1;32m    984\u001b[0m   \u001b[38;5;66;03m# constraint.\u001b[39;00m\n",
            "\u001b[0;31mValueError\u001b[0m: compiling computation that requires 10000000 logical devices, but only 4 XLA devices are available (num_replicas=10000000)"
          ]
        }
      ],
      "source": [
        "x_pmap = dot_parallel(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 153,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1DCjJciai7Z_",
        "outputId": "be77bca5-7898-4418-838b-e0dc336fd0d3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10000000, 3)"
            ]
          },
          "execution_count": 153,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 154,
      "metadata": {
        "id": "0zaUwp4tjroC"
      },
      "outputs": [],
      "source": [
        "v1sp = v1s.reshape((2, v1s.shape[0]//2, v1s.shape[1]))\n",
        "v2sp = v2s.reshape((2, v2s.shape[0]//2, v2s.shape[1]))    # we use integer division here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 155,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ntlAFzYAi9_s",
        "outputId": "04ae1775-877d-4746-f5a4-8a0cd503ec89"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(2, 5000000, 3)"
            ]
          },
          "execution_count": 155,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1sp.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 156,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ttr1LezXIvos",
        "outputId": "56152362-0aae-4dde-b8b7-40943627c53e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(Array([1.456546  , 1.1449531 , 0.02485494], dtype=float32),\n",
              " Array([1.456546  , 1.1449531 , 0.02485494], dtype=float32))"
            ]
          },
          "execution_count": 156,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s[0,:], v1sp[0,0,:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 157,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xQJAXG1MI6ja",
        "outputId": "e26e2375-5d53-4bf3-c395-3b6f7cddd297"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(Array([ 0.50284576,  0.38336843, -0.35499337], dtype=float32),\n",
              " Array([-0.3869151 ,  0.31480923, -0.11581508], dtype=float32))"
            ]
          },
          "execution_count": 157,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s[1250000,:], v1sp[1,0,:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 158,
      "metadata": {
        "id": "kED5PhIvj_Yz"
      },
      "outputs": [],
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 159,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C4SjsR23KaYm",
        "outputId": "b8396ca5-4be1-449c-ecea-ba5288bc266d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(2,)"
            ]
          },
          "execution_count": 159,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 160,
      "metadata": {
        "id": "-5sG0zdjkNwR"
      },
      "outputs": [],
      "source": [
        "dot_parallel = jax.pmap(jax.vmap(dot))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 161,
      "metadata": {
        "id": "gCI71XoukYOv"
      },
      "outputs": [],
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 162,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W4oj2g-bkgbM",
        "outputId": "e437750a-d162-4749-af79-4ab67c98bc1f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(2, 5000000)"
            ]
          },
          "execution_count": 162,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dyIOI3cakjL0",
        "outputId": "0e0cebc0-8e86-4961-fbe0-9ef3ccc50936"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "jaxlib.xla_extension.ArrayImpl"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(x_pmap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1W2vIrDxklDX",
        "outputId": "cff17983-2c42-4683-a47a-346921b3392c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "38UCSdMLmPYS",
        "outputId": "00a13cde-c89c-4206-dc20-ec650e88f086"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.numpy.all(x_pmap == x_vmap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k2jjZe-mo47n",
        "outputId": "76ab37a2-2035-4a8d-e981-419ead1dcb10"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "5.68 ms ± 3.3 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit xp = dot_parallel(v1sp,v2sp).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-_jLRzyBRDTJ",
        "outputId": "9d5e09f6-e779-475d-be17-b812d0aab908"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.1981742 , -3.0097814 , -2.539115  , ..., -2.7281013 ,\n",
              "        0.26364914,  1.9583547 ], dtype=float32)"
            ]
          },
          "execution_count": 33,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TxECifZoRGYY",
        "outputId": "ca2afae7-9e4a-4238-9cb8-14e669c2e680"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.1981742 , -3.0097814 , -2.539115  , ..., -2.7281013 ,\n",
              "        0.26364914,  1.9583547 ], dtype=float32)"
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_vmap"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ReSRcPnlR_7g"
      },
      "source": [
        "A very simple example when you can easily switch vmap and pmap:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "R28QzddqW6Xm"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "giEtn52mSIk-"
      },
      "outputs": [],
      "source": [
        "vs = random.normal(rng_key, shape=(8,3))\n",
        "v1s = vs[:4,:]\n",
        "v2s = vs[4:,:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LfT1rhKJTgwq",
        "outputId": "63bb34b2-4461-4bfe-d853-bea782581fbf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 40,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.vmap(dot)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z46RrWmhTnEK",
        "outputId": "5972ad49-cb7b-4d10-9a2c-7d453c675035"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 41,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.pmap(dot)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "id": "vNVXUxJpXLaB"
      },
      "outputs": [],
      "source": [
        "dot_v = jax.jit(jax.vmap(dot))\n",
        "x = dot_v(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TbZv85rgXRR_",
        "outputId": "0bd2420a-91d6-4299-844e-7d893344c667"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/jetjiang/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:1826: UserWarning: The jitted function dot includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See https://github.com/google/jax/issues/2926.\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "dot_pjo = jax.jit(jax.pmap(dot))\n",
        "x = dot_pjo(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "id": "OHXNHJTYYJb5"
      },
      "outputs": [],
      "source": [
        "dot_pji = jax.pmap(jax.jit(dot))\n",
        "x = dot_pji(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "KNrEftwLXZlw"
      },
      "outputs": [],
      "source": [
        "dot_p = jax.pmap(dot)\n",
        "x = dot_p(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pm9Q-E4UWcHs",
        "outputId": "56228848-94bf-4e56-bf3e-eb024ad4ccb0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "116 μs ± 446 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit dot_v(v1s,v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IJtuPjyBWiBp",
        "outputId": "c1e1e811-cb0f-4624-cb1e-0400270f6c50"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "487 μs ± 10.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit dot_pjo(v1s,v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q6q8J8TCYQho",
        "outputId": "a3be9b08-b156-44c9-afb0-793a5789a711"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "841 μs ± 3.65 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit dot_pji(v1s,v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KWDwONE-Xsjn",
        "outputId": "13ff8763-b842-464a-d4e6-85f1a6a87e9e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "827 μs ± 9.88 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit dot_p(v1s,v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lJFuKu1XAhFj",
        "outputId": "7c2229a5-0767-44af-89b1-7aa8f6f7d0e0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4,3] b:f32[4,3]. let\n",
              "    c:f32[4] = pjit[\n",
              "      name=dot\n",
              "      jaxpr={ lambda ; d:f32[4,3] e:f32[4,3]. let\n",
              "          f:f32[4] = dot_general[\n",
              "            dimension_numbers=(([1], [1]), ([0], [0]))\n",
              "            preferred_element_type=float32\n",
              "          ] d e\n",
              "        in (f,) }\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 50,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(dot_v)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PJE-4HWRAnko",
        "outputId": "c60f2960-ac34-4733-f137-5f53f57da97a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4,3] b:f32[4,3]. let\n",
              "    c:f32[4] = pjit[\n",
              "      name=dot\n",
              "      jaxpr={ lambda ; d:f32[4,3] e:f32[4,3]. let\n",
              "          f:f32[4] = xla_pmap[\n",
              "            axis_name=<axis 0x7fc329dead40>\n",
              "            axis_size=4\n",
              "            backend=None\n",
              "            call_jaxpr={ lambda ; g:f32[3] h:f32[3]. let\n",
              "                i:f32[] = dot_general[\n",
              "                  dimension_numbers=(([0], [0]), ([], []))\n",
              "                  preferred_element_type=float32\n",
              "                ] g h\n",
              "              in (i,) }\n",
              "            devices=None\n",
              "            donated_invars=(False, False)\n",
              "            global_axis_size=4\n",
              "            in_axes=(0, 0)\n",
              "            is_explicit_global_axis_size=False\n",
              "            name=dot\n",
              "            out_axes=(0,)\n",
              "          ] d e\n",
              "        in (f,) }\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 51,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(dot_pjo)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VH60u2geApK3",
        "outputId": "6881a604-42cb-44bf-b27e-e52d8c87f1ca"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4,3] b:f32[4,3]. let\n",
              "    c:f32[4] = xla_pmap[\n",
              "      axis_name=<axis 0x7fd8586c3d30>\n",
              "      axis_size=4\n",
              "      backend=None\n",
              "      call_jaxpr={ lambda ; d:f32[3] e:f32[3]. let\n",
              "          f:f32[] = pjit[\n",
              "            name=dot\n",
              "            jaxpr={ lambda ; g:f32[3] h:f32[3]. let\n",
              "                i:f32[] = dot_general[\n",
              "                  dimension_numbers=(([0], [0]), ([], []))\n",
              "                  preferred_element_type=float32\n",
              "                ] g h\n",
              "              in (i,) }\n",
              "          ] d e\n",
              "        in (f,) }\n",
              "      devices=None\n",
              "      donated_invars=(False, False)\n",
              "      global_axis_size=4\n",
              "      in_axes=(0, 0)\n",
              "      is_explicit_global_axis_size=False\n",
              "      name=dot\n",
              "      out_axes=(0,)\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 52,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(dot_pji)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Alc3_10BAotr",
        "outputId": "7561b326-de10-4f1a-a264-06085045a764"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4,3] b:f32[4,3]. let\n",
              "    c:f32[4] = xla_pmap[\n",
              "      axis_name=<axis 0x7fc329dead40>\n",
              "      axis_size=4\n",
              "      backend=None\n",
              "      call_jaxpr={ lambda ; d:f32[3] e:f32[3]. let\n",
              "          f:f32[] = dot_general[\n",
              "            dimension_numbers=(([0], [0]), ([], []))\n",
              "            preferred_element_type=float32\n",
              "          ] d e\n",
              "        in (f,) }\n",
              "      devices=None\n",
              "      donated_invars=(False, False)\n",
              "      global_axis_size=4\n",
              "      in_axes=(0, 0)\n",
              "      is_explicit_global_axis_size=False\n",
              "      name=dot\n",
              "      out_axes=(0,)\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(dot_p)(v1s,v2s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rB5gg9rvePQb"
      },
      "source": [
        "## Controlling pmap() behavior"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1_dx8YVoGb-H"
      },
      "source": [
        "### Using in_axes parameter"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82lZ7bqNT4TN"
      },
      "source": [
        "Using a small array just for demonstration purposes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "WKUM9E0jT3sx"
      },
      "outputs": [],
      "source": [
        "vs = random.normal(rng_key, shape=(8,3))\n",
        "v1s = vs[:4,:]\n",
        "v2s = vs[4:,:]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "pmydsax3UEkA"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YxGyiGP5mlcK"
      },
      "source": [
        "A default value:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "id": "_anvIZEhj-bk"
      },
      "outputs": [],
      "source": [
        "dot_pmapped = jax.pmap(dot, in_axes=(0,0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZjEzc-wedMDq",
        "outputId": "5791ec39-40b3-4839-ef9e-9b5c056fd850"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 57,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dot_pmapped(v1s, v2s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvjNEkoUN-l"
      },
      "source": [
        "What if one of our arrays is transposed (the mapping axis is the second one)?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3CNhCLoZal3p",
        "outputId": "e08db0a7-c13a-4c4a-bb1a-cef61c2c75dc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((3, 4), (4, 3))"
            ]
          },
          "execution_count": 58,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.T.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jr-GZ4zmUNXa",
        "outputId": "21f6a200-9d6f-4d13-acfb-e0f1a40a90fd"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 59,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.pmap(dot, in_axes=(1,0))(v1s.T, v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 332
        },
        "id": "4ZTIPUEyUeqJ",
        "outputId": "7a0b1619-820c-4dd9-dae7-07e4fa0a44d5"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "pmap got inconsistent sizes for array axes to be mapped:\n  * one axis had size 3: axis 0 of argument v1 of type float32[3,4];\n  * one axis had size 4: axis 0 of argument v2 of type float32[4,3]",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[60], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mdot_pmapped\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1s\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mT\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2s\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 3 frame]\u001b[0m\n",
            "File \u001b[0;32m~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/api.py:1315\u001b[0m, in \u001b[0;36m_mapped_axis_size\u001b[0;34m(fn, tree, vals, dims, name)\u001b[0m\n\u001b[1;32m   1313\u001b[0m   \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1314\u001b[0m     msg\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m  * some axes (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mct\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m of them) had size \u001b[39m\u001b[38;5;132;01m{\u001b[39;00msz\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, e.g. axis \u001b[39m\u001b[38;5;132;01m{\u001b[39;00max\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mex\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m;\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 1315\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(msg)[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m2\u001b[39m])\n",
            "\u001b[0;31mValueError\u001b[0m: pmap got inconsistent sizes for array axes to be mapped:\n  * one axis had size 3: axis 0 of argument v1 of type float32[3,4];\n  * one axis had size 4: axis 0 of argument v2 of type float32[4,3]"
          ]
        }
      ],
      "source": [
        "dot_pmapped(v1s.T, v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QQub20ZTauM9",
        "outputId": "729a6a83-f3e3-42a6-8254-13005140337d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((3, 4), (3, 4))"
            ]
          },
          "execution_count": 61,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.T.shape, v2s.T.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HXY3CZ9uUjIb",
        "outputId": "fb81fc11-cdda-46a1-f31b-1161cbf3a51f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 62,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.pmap(dot, in_axes=(1,1))(v1s.T, v2s.T)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rgsH99-1UoWA",
        "outputId": "5fb5baa3-d7cf-4680-b264-307542483174"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([-3.3310347,  3.7464314,  0.7410331], dtype=float32)"
            ]
          },
          "execution_count": 63,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dot_pmapped(v1s.T, v2s.T)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gqz7L6uSmqPc"
      },
      "source": [
        "A more complicated case:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {
        "id": "q7ZB2g1GdWJz"
      },
      "outputs": [],
      "source": [
        "def scaled_dot(v1, v2, koeff):\n",
        "  return koeff*jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {
        "id": "yjHKGd1-diSw"
      },
      "outputs": [],
      "source": [
        "v1s_ = v1s\n",
        "v2s_ = v2s.T\n",
        "k = 1.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LU0SfSJieggS",
        "outputId": "43d5171f-8a50-4e0e-f2c2-aa61dcd9d1eb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((4, 3), (3, 4))"
            ]
          },
          "execution_count": 66,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s_.shape, v2s_.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {
        "id": "ymOex9oHemQ8"
      },
      "outputs": [],
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q-pqO8b_muy4"
      },
      "source": [
        "Default values do not work:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 523
        },
        "id": "dbFc89nYfaH8",
        "outputId": "22995d8e-1604-43b3-d6df-f2851afeccd0"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "pmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "File \u001b[0;32m~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/api.py:1253\u001b[0m, in \u001b[0;36m_mapped_axis_size.<locals>._get_axis_size\u001b[0;34m(name, shape, axis)\u001b[0m\n\u001b[1;32m   1252\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1253\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mshape\u001b[49m\u001b[43m[\u001b[49m\u001b[43maxis\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m   1254\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mIndexError\u001b[39;00m, \u001b[38;5;167;01mTypeError\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m e:\n",
            "\u001b[0;31mIndexError\u001b[0m: tuple index out of range",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[68], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mscaled_dot_pmapped\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1s_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2s_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
            "File \u001b[0;32m~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/api.py:1257\u001b[0m, in \u001b[0;36m_mapped_axis_size.<locals>._get_axis_size\u001b[0;34m(name, shape, axis)\u001b[0m\n\u001b[1;32m   1255\u001b[0m min_rank \u001b[38;5;241m=\u001b[39m axis \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m axis \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;241m-\u001b[39maxis\n\u001b[1;32m   1256\u001b[0m \u001b[38;5;66;03m# TODO(mattjj): better error message here\u001b[39;00m\n\u001b[0;32m-> 1257\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   1258\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m was requested to map its argument along axis \u001b[39m\u001b[38;5;132;01m{\u001b[39;00maxis\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1259\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwhich implies that its rank should be at least \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmin_rank\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1260\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbut is only \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(shape)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m (its shape is \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n",
            "\u001b[0;31mValueError\u001b[0m: pmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())"
          ]
        }
      ],
      "source": [
        "scaled_dot_pmapped(v1s_, v2s_, k)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {
        "id": "9dTFcAHQfdVf"
      },
      "outputs": [],
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot, in_axes=(0,1,None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A2f452A9flEO",
        "outputId": "0522f2ee-a1b2-4940-d05d-1a5a34cb06bb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 70,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "scaled_dot_pmapped(v1s_, v2s_, k)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PVSBhIYAmyi4"
      },
      "source": [
        "Using more complex parameter structure:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "id": "6cuNZDUZiMEH"
      },
      "outputs": [],
      "source": [
        "def scaled_dot(data, koeff):\n",
        "  return koeff*jnp.vdot(data['a'], data['b'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "id": "pfdMJ7eeijRq"
      },
      "outputs": [],
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot, in_axes=({'a':0,'b':1},None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lMjD-rsCivBv",
        "outputId": "1441d7fc-62c7-4223-c49f-1c09708f8362"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.4808662 ,  0.2834839 ,  0.28449908, -0.8924192 ], dtype=float32)"
            ]
          },
          "execution_count": 73,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "scaled_dot_pmapped({'a':v1s_, 'b': v2s_}, k)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72VWPGXrm5du"
      },
      "source": [
        "### Using out_axes parameter"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "id": "dv3wsrlJm7BB"
      },
      "outputs": [],
      "source": [
        "def scale(v, koeff):\n",
        "  return koeff*v"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "id": "C34aK9bgnBed"
      },
      "outputs": [],
      "source": [
        "scale_pmapped = jax.pmap(scale,\n",
        "                         in_axes=(0,None),\n",
        "                         out_axes=(1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "id": "DodCkW4oiz2p"
      },
      "outputs": [],
      "source": [
        "res = scale_pmapped(v1s, 2.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0UTVvz2ZgPQj",
        "outputId": "79ac19dc-00cb-4121-c137-dcfe0f00a476"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((4, 3), (3, 4))"
            ]
          },
          "execution_count": 77,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.shape, res.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "DGrRn25qnYrj"
      },
      "outputs": [],
      "source": [
        "scale_pmapped = jax.pmap(scale, in_axes=(0,None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4qzPVn_znZ3o",
        "outputId": "95220e2d-757a-4bf3-a5e9-048dfbd04d74"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[ 1.1957493 ,  2.100962  ,  4.260094  ],\n",
              "       [-2.7670264 , -3.6707642 , -0.10890777],\n",
              "       [-1.8921121 ,  0.99184775,  0.52138793],\n",
              "       [ 1.7561268 ,  2.1729617 ,  2.6528773 ]], dtype=float32)"
            ]
          },
          "execution_count": 79,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "scale_pmapped(v1s, 2.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "id": "K5Jps90J6HLw"
      },
      "outputs": [],
      "source": [
        "scale_pmapped = jax.pmap(scale,\n",
        "                         in_axes=(0,None),\n",
        "                         out_axes=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ru8LmRBtYavl"
      },
      "source": [
        "## A more real-life case with larger tensors and mixing vmap and pmap"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "id": "wGaZtjJ1hqw7"
      },
      "outputs": [],
      "source": [
        "vs = random.normal(rng_key, shape=(20_000_000,3))\n",
        "v1s = vs[:10_000_000,:].T\n",
        "v2s = vs[10_000_000:,:].T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 82,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hasxnv11h-bb",
        "outputId": "85d5d103-ddcc-48b4-9e42-ba1d04254cba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((3, 10000000), (3, 10000000))"
            ]
          },
          "execution_count": 82,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1s.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lOkDY2wromwx",
        "outputId": "ad86e702-3b3d-4131-8972-3391a08181d1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((3, 4, 2500000), (3, 4, 2500000))"
            ]
          },
          "execution_count": 83,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "v1sp = v1s.reshape((v1s.shape[0], 4, v1s.shape[1]//4))\n",
        "v2sp = v2s.reshape((v2s.shape[0], 4, v2s.shape[1]//4))\n",
        "v1sp.shape, v2sp.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "id": "wRgFXKt-o0aQ"
      },
      "outputs": [],
      "source": [
        "dot_parallel = jax.pmap(\n",
        "    jax.vmap(dot, in_axes=(1,1)),\n",
        "    in_axes=(1,1)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 85,
      "metadata": {
        "id": "tEa0RhW0pW5X"
      },
      "outputs": [],
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wxz83U68pb0I",
        "outputId": "bd89f3c7-c246-4d24-cada-cb9fbae433a5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 2500000)"
            ]
          },
          "execution_count": 86,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 87,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1kEHvzk5pmLe",
        "outputId": "8f200aed-ae8e-4e8d-d1cb-57112a5dc17d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "execution_count": 87,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m-AaCUnGpp4l",
        "outputId": "d89f63b7-2d1b-4514-ccb2-58822b37da22"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "execution_count": 88,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.numpy.all(x_pmap == x_vmap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yd_Y1-YOpvdU",
        "outputId": "6d9b62da-9cfe-4533-8792-1188bb0a778d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.1981742, -3.0097814, -2.539115 ], dtype=float32)"
            ]
          },
          "execution_count": 89,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_pmap[:3]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0VbvxplMpy22",
        "outputId": "ae25a93b-7279-4f71-c983-8c2954637c3b"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 1.1981742, -3.0097814, -2.539115 ], dtype=float32)"
            ]
          },
          "execution_count": 90,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x_vmap[:3]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7zxdlwUP14or"
      },
      "source": [
        "## Using collective ops and the axis_name parameter"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dF2ONW2AhMVd"
      },
      "source": [
        "### Normalization examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "POmEA_DwhRD6"
      },
      "source": [
        "A small array example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {
        "id": "DrWA7KPZ196e"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array(range(4))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iGDJGEfn245M",
        "outputId": "df778b2d-75b2-4f49-b829-20c8bc18c460"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([0, 1, 2, 3], dtype=int32)"
            ]
          },
          "execution_count": 96,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {
        "id": "AYLt3IU8_pFE"
      },
      "outputs": [],
      "source": [
        "norm = jax.pmap(\n",
        "    lambda x: x/jax.lax.psum(x, axis_name='p'),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2KXDuNJi_zv7",
        "outputId": "fe8f62ce-4762-4f11-f5ca-2f0c0b3627d5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([0.        , 0.16666667, 0.33333334, 0.5       ], dtype=float32)"
            ]
          },
          "execution_count": 98,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "norm(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5VzX8Gc4_1vo",
        "outputId": "1f172ea5-875d-4e0e-9a97-798311962101"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "execution_count": 99,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(norm(arr))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BAsivlcghTHm"
      },
      "source": [
        "A large array (more elements than XLA devices) example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "id": "MEDDfwp21SH8"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array(range(200))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 101,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4otHYX5R1VKR",
        "outputId": "392dd228-748e-436c-de59-9c2bfe1f9232"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 50)"
            ]
          },
          "execution_count": 101,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr = arr.reshape(4, 50)\n",
        "arr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 102,
      "metadata": {
        "id": "4n21Oyvp1Z-t"
      },
      "outputs": [],
      "source": [
        "norm = jax.pmap(\n",
        "    lambda x: x/jax.lax.psum(jnp.sum(x), axis_name='p'),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 103,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fCHy2q-n1eah",
        "outputId": "678f2166-ad1b-4bbc-a01e-bed5a1e9547c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 50)"
            ]
          },
          "execution_count": 103,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SUEfvESE4bn_",
        "outputId": "8bf0afd3-678b-4a54-8664-38e5f782881e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "execution_count": 104,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(narr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0lFI5h1XhyqI"
      },
      "source": [
        "Using groups"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 107,
      "metadata": {
        "id": "y3Zt5UpD-l3L"
      },
      "outputs": [],
      "source": [
        "norm = jax.pmap(\n",
        "    lambda x: x/jax.lax.psum(\n",
        "        jnp.sum(x),\n",
        "        axis_name='p',\n",
        "        axis_index_groups=[[0,1], [2,3]]\n",
        "    ),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 108,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "58qiMx2e_E8l",
        "outputId": "730d5df7-69f4-4548-b16f-5d54178d7673"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 50)"
            ]
          },
          "execution_count": 108,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 109,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iDxWZ8xN_VqZ",
        "outputId": "01ad325a-31e1-4621-92a0-d7e205030567"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(2.0000002, dtype=float32)"
            ]
          },
          "execution_count": 109,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(narr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 110,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zTRiyInw_lmT",
        "outputId": "600d2d59-5c86-49cb-ba1f-8beaff122cc9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(Array(1., dtype=float32), Array(1., dtype=float32))"
            ]
          },
          "execution_count": 110,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(narr[:2]), jnp.sum(narr[2:4])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mAFZXq76DQbe"
      },
      "source": [
        "### Using nested pmap() and vmap()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 111,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xY7giT0-DP_C",
        "outputId": "15c11967-83aa-472b-97b7-e7627735b582"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 50)"
            ]
          },
          "execution_count": 111,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr = jnp.array(range(200))\n",
        "arr = arr.reshape(4, 50)\n",
        "arr.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uNs_Sri4hf7c"
      },
      "source": [
        "Understanding different axes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 112,
      "metadata": {
        "id": "S54zj0qMFxNb"
      },
      "outputs": [],
      "source": [
        "f = jax.pmap(\n",
        "    jax.vmap(\n",
        "        lambda x: jax.lax.pmax(x, axis_name='p'),    # finding a maximum across pmax axis\n",
        "        axis_name='v'\n",
        "    ),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 113,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZiMDUbLlFzl8",
        "outputId": "ba96c985-7a70-4685-d9e0-228d010e25e1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162,\n",
              "        163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,\n",
              "        176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188,\n",
              "        189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162,\n",
              "        163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,\n",
              "        176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188,\n",
              "        189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162,\n",
              "        163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,\n",
              "        176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188,\n",
              "        189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162,\n",
              "        163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175,\n",
              "        176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188,\n",
              "        189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199]],      dtype=int32)"
            ]
          },
          "execution_count": 113,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "f(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 120,
      "metadata": {
        "id": "47d7KE6PfwqW"
      },
      "outputs": [],
      "source": [
        "f = jax.pmap(\n",
        "    jax.vmap(\n",
        "        lambda x: jax.lax.pmax(x, axis_name=('p', 'v')),    # Finding global maximum across two axes\n",
        "        axis_name='v'\n",
        "    ),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 121,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-blZXHawf0sp",
        "outputId": "822ebd77-c216-42b1-b6f9-d404bbbe4526"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199],\n",
              "       [199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199],\n",
              "       [199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199],\n",
              "       [199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199,\n",
              "        199, 199, 199, 199, 199, 199, 199, 199, 199, 199, 199]],      dtype=int32)"
            ]
          },
          "execution_count": 121,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "f(arr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DfLX3pMphlg1"
      },
      "source": [
        "Using both axes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 116,
      "metadata": {
        "id": "WhVL-R8qDd7_"
      },
      "outputs": [],
      "source": [
        "f = jax.pmap(\n",
        "    jax.vmap(\n",
        "        lambda x: jax.lax.pmax(x, axis_name='v')/jax.lax.pmax(x, axis_name='p'),\n",
        "        axis_name='v'\n",
        "    ),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 117,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rpNlgksOEEpJ",
        "outputId": "214388ee-096c-4b3a-ca45-f5e61f0f1b4a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[0.32666668, 0.3245033 , 0.3223684 , 0.32026145, 0.31818178,\n",
              "        0.31612903, 0.31410256, 0.3121019 , 0.31012657, 0.3081761 ,\n",
              "        0.30625   , 0.3043478 , 0.30246916, 0.30061352, 0.29878047,\n",
              "        0.29696968, 0.2951807 , 0.2934132 , 0.2916667 , 0.28994083,\n",
              "        0.2882353 , 0.28654972, 0.2848837 , 0.28323698, 0.2816092 ,\n",
              "        0.28      , 0.2784091 , 0.27683613, 0.2752809 , 0.273743  ,\n",
              "        0.27222222, 0.27071825, 0.26923078, 0.2677596 , 0.26630434,\n",
              "        0.26486486, 0.26344085, 0.2620321 , 0.2606383 , 0.25925925,\n",
              "        0.25789472, 0.2565445 , 0.25520834, 0.253886  , 0.2525773 ,\n",
              "        0.25128207, 0.25000003, 0.24873096, 0.24747474, 0.24623115],\n",
              "       [0.66      , 0.65562916, 0.6513158 , 0.64705884, 0.6428571 ,\n",
              "        0.63870966, 0.6346154 , 0.6305733 , 0.62658226, 0.6226415 ,\n",
              "        0.61875004, 0.6149068 , 0.61111116, 0.607362  , 0.6036585 ,\n",
              "        0.59999996, 0.59638554, 0.5928144 , 0.58928573, 0.5857988 ,\n",
              "        0.58235294, 0.57894737, 0.5755814 , 0.5722543 , 0.5689655 ,\n",
              "        0.5657143 , 0.5625    , 0.559322  , 0.55617976, 0.55307263,\n",
              "        0.55      , 0.54696137, 0.54395604, 0.5409836 , 0.5380435 ,\n",
              "        0.5351351 , 0.5322581 , 0.5294118 , 0.5265957 , 0.5238095 ,\n",
              "        0.5210526 , 0.5183246 , 0.515625  , 0.51295334, 0.5103093 ,\n",
              "        0.50769234, 0.5051021 , 0.5025381 , 0.5       , 0.49748743],\n",
              "       [0.99333334, 0.98675495, 0.9802632 , 0.9738562 , 0.9675324 ,\n",
              "        0.9612903 , 0.9551282 , 0.9490446 , 0.9430379 , 0.9371069 ,\n",
              "        0.93125004, 0.92546576, 0.91975313, 0.9141105 , 0.90853655,\n",
              "        0.9030303 , 0.89759034, 0.8922156 , 0.8869048 , 0.8816568 ,\n",
              "        0.8764706 , 0.87134504, 0.86627907, 0.8612717 , 0.8563218 ,\n",
              "        0.85142857, 0.84659094, 0.84180784, 0.83707863, 0.8324022 ,\n",
              "        0.8277778 , 0.82320446, 0.81868136, 0.8142077 , 0.8097826 ,\n",
              "        0.8054054 , 0.8010753 , 0.79679143, 0.7925532 , 0.78835976,\n",
              "        0.7842105 , 0.7801047 , 0.7760417 , 0.7720207 , 0.7680412 ,\n",
              "        0.7641026 , 0.76020414, 0.75634515, 0.75252527, 0.7487437 ],\n",
              "       [1.3266667 , 1.3178807 , 1.3092105 , 1.3006536 , 1.2922077 ,\n",
              "        1.2838709 , 1.2756411 , 1.2675159 , 1.2594936 , 1.2515723 ,\n",
              "        1.24375   , 1.2360247 , 1.2283951 , 1.2208589 , 1.2134145 ,\n",
              "        1.2060605 , 1.1987951 , 1.1916168 , 1.1845238 , 1.1775148 ,\n",
              "        1.1705883 , 1.1637427 , 1.1569767 , 1.1502889 , 1.1436782 ,\n",
              "        1.1371429 , 1.1306819 , 1.1242937 , 1.1179775 , 1.1117318 ,\n",
              "        1.1055555 , 1.0994476 , 1.0934067 , 1.0874318 , 1.0815217 ,\n",
              "        1.0756756 , 1.0698925 , 1.0641712 , 1.0585107 , 1.05291   ,\n",
              "        1.0473684 , 1.0418848 , 1.0364584 , 1.0310881 , 1.0257732 ,\n",
              "        1.0205128 , 1.0153062 , 1.0101522 , 1.0050505 , 1.        ]],      dtype=float32)"
            ]
          },
          "execution_count": 117,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "f(arr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D4o8QRgXhoJP"
      },
      "source": [
        "Global normalization using two axes at once"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 122,
      "metadata": {
        "id": "eqZ8GBcFfttm"
      },
      "outputs": [],
      "source": [
        "norm = jax.pmap(\n",
        "    jax.vmap(\n",
        "        lambda x: x/jax.lax.psum(x, axis_name=('p','v')),\n",
        "        axis_name='v'\n",
        "    ),\n",
        "    axis_name='p')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 123,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GKZKJTyKiBCy",
        "outputId": "2af26944-96f2-4cb5-e5df-4f6868e47ccb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4, 50)"
            ]
          },
          "execution_count": 123,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 124,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Weeis4uNiGmI",
        "outputId": "f450ba48-7042-441e-8a7e-52d2cbb2cdf8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "execution_count": 124,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(narr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 125,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NIc0LfUeiON_",
        "outputId": "3cecb9f0-dfbe-4731-915f-db354e0bc1f3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[0.0000000e+00, 5.0251256e-05, 1.0050251e-04, 1.5075377e-04,\n",
              "        2.0100502e-04, 2.5125628e-04, 3.0150753e-04, 3.5175879e-04,\n",
              "        4.0201005e-04, 4.5226130e-04, 5.0251256e-04, 5.5276381e-04,\n",
              "        6.0301507e-04, 6.5326632e-04, 7.0351758e-04, 7.5376884e-04,\n",
              "        8.0402009e-04, 8.5427135e-04, 9.0452260e-04, 9.5477386e-04,\n",
              "        1.0050251e-03, 1.0552764e-03, 1.1055276e-03, 1.1557789e-03,\n",
              "        1.2060301e-03, 1.2562814e-03, 1.3065326e-03, 1.3567839e-03,\n",
              "        1.4070352e-03, 1.4572864e-03, 1.5075377e-03, 1.5577889e-03,\n",
              "        1.6080402e-03, 1.6582914e-03, 1.7085427e-03, 1.7587940e-03,\n",
              "        1.8090452e-03, 1.8592965e-03, 1.9095477e-03, 1.9597989e-03,\n",
              "        2.0100502e-03, 2.0603016e-03, 2.1105527e-03, 2.1608039e-03,\n",
              "        2.2110553e-03, 2.2613066e-03, 2.3115578e-03, 2.3618089e-03,\n",
              "        2.4120603e-03, 2.4623116e-03],\n",
              "       [2.5125628e-03, 2.5628139e-03, 2.6130653e-03, 2.6633167e-03,\n",
              "        2.7135678e-03, 2.7638189e-03, 2.8140703e-03, 2.8643217e-03,\n",
              "        2.9145728e-03, 2.9648240e-03, 3.0150753e-03, 3.0653267e-03,\n",
              "        3.1155779e-03, 3.1658290e-03, 3.2160804e-03, 3.2663317e-03,\n",
              "        3.3165829e-03, 3.3668340e-03, 3.4170854e-03, 3.4673368e-03,\n",
              "        3.5175879e-03, 3.5678390e-03, 3.6180904e-03, 3.6683418e-03,\n",
              "        3.7185929e-03, 3.7688441e-03, 3.8190954e-03, 3.8693468e-03,\n",
              "        3.9195977e-03, 3.9698491e-03, 4.0201005e-03, 4.0703518e-03,\n",
              "        4.1206032e-03, 4.1708541e-03, 4.2211055e-03, 4.2713569e-03,\n",
              "        4.3216078e-03, 4.3718591e-03, 4.4221105e-03, 4.4723619e-03,\n",
              "        4.5226132e-03, 4.5728642e-03, 4.6231155e-03, 4.6733669e-03,\n",
              "        4.7236178e-03, 4.7738692e-03, 4.8241206e-03, 4.8743719e-03,\n",
              "        4.9246233e-03, 4.9748742e-03],\n",
              "       [5.0251256e-03, 5.0753769e-03, 5.1256279e-03, 5.1758792e-03,\n",
              "        5.2261306e-03, 5.2763820e-03, 5.3266333e-03, 5.3768842e-03,\n",
              "        5.4271356e-03, 5.4773870e-03, 5.5276379e-03, 5.5778893e-03,\n",
              "        5.6281406e-03, 5.6783920e-03, 5.7286434e-03, 5.7788943e-03,\n",
              "        5.8291457e-03, 5.8793970e-03, 5.9296479e-03, 5.9798993e-03,\n",
              "        6.0301507e-03, 6.0804021e-03, 6.1306534e-03, 6.1809043e-03,\n",
              "        6.2311557e-03, 6.2814071e-03, 6.3316580e-03, 6.3819094e-03,\n",
              "        6.4321607e-03, 6.4824121e-03, 6.5326635e-03, 6.5829144e-03,\n",
              "        6.6331658e-03, 6.6834171e-03, 6.7336680e-03, 6.7839194e-03,\n",
              "        6.8341708e-03, 6.8844222e-03, 6.9346735e-03, 6.9849244e-03,\n",
              "        7.0351758e-03, 7.0854272e-03, 7.1356781e-03, 7.1859295e-03,\n",
              "        7.2361808e-03, 7.2864322e-03, 7.3366836e-03, 7.3869345e-03,\n",
              "        7.4371858e-03, 7.4874372e-03],\n",
              "       [7.5376881e-03, 7.5879395e-03, 7.6381909e-03, 7.6884422e-03,\n",
              "        7.7386936e-03, 7.7889445e-03, 7.8391954e-03, 7.8894468e-03,\n",
              "        7.9396982e-03, 7.9899495e-03, 8.0402009e-03, 8.0904523e-03,\n",
              "        8.1407037e-03, 8.1909550e-03, 8.2412064e-03, 8.2914568e-03,\n",
              "        8.3417082e-03, 8.3919596e-03, 8.4422110e-03, 8.4924623e-03,\n",
              "        8.5427137e-03, 8.5929651e-03, 8.6432155e-03, 8.6934669e-03,\n",
              "        8.7437183e-03, 8.7939696e-03, 8.8442210e-03, 8.8944724e-03,\n",
              "        8.9447238e-03, 8.9949751e-03, 9.0452265e-03, 9.0954769e-03,\n",
              "        9.1457283e-03, 9.1959797e-03, 9.2462311e-03, 9.2964824e-03,\n",
              "        9.3467338e-03, 9.3969852e-03, 9.4472356e-03, 9.4974870e-03,\n",
              "        9.5477384e-03, 9.5979897e-03, 9.6482411e-03, 9.6984925e-03,\n",
              "        9.7487438e-03, 9.7989952e-03, 9.8492466e-03, 9.8994970e-03,\n",
              "        9.9497484e-03, 9.9999998e-03]], dtype=float32)"
            ]
          },
          "execution_count": 125,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "narr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cs_oFSYjkaIN"
      },
      "source": [
        "Nested pmap"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 127,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W-dMX9EckckK",
        "outputId": "0227f267-32c5-4786-91e7-c68955e12492"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[0, 1],\n",
              "       [2, 3]], dtype=int32)"
            ]
          },
          "execution_count": 127,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr = jnp.array(range(4)).reshape(2,2)\n",
        "arr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 128,
      "metadata": {
        "id": "YVEFRvlYkvcW"
      },
      "outputs": [],
      "source": [
        "n = jax.pmap(\n",
        "    jax.pmap(\n",
        "        lambda x: x/jax.lax.psum(x, axis_name=('rows','cols')),\n",
        "        axis_name='cols'\n",
        "    ),\n",
        "    axis_name='rows')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 129,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LVcHVBl9k7Tg",
        "outputId": "de33ffae-10e9-4c64-dd7e-ef0d4decfdee"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([[0.        , 0.16666667],\n",
              "       [0.33333334, 0.5       ]], dtype=float32)"
            ]
          },
          "execution_count": 129,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "n(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 130,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gAVypn8Lk-Ov",
        "outputId": "378a2b68-e69b-4441-bfb0-ae035fe303af"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "execution_count": 130,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(n(arr))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vY_seo4knd9M"
      },
      "source": [
        "The same example using the decorator style"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 163,
      "metadata": {
        "id": "RiikZf-lni6V"
      },
      "outputs": [],
      "source": [
        "from functools import partial"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 164,
      "metadata": {
        "id": "uBtUfQgonjer"
      },
      "outputs": [],
      "source": [
        "@partial(jax.pmap, axis_name='rows')\n",
        "@partial(jax.pmap, axis_name='cols')\n",
        "def n(x):\n",
        "  return x/jax.lax.psum(x, axis_name=('rows','cols'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 165,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XQeCcTBxnwdu",
        "outputId": "92353719-280a-4e74-fcd6-85f7a798efc7"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "execution_count": 165,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(n(arr))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DVXNIzeGNkaF"
      },
      "source": [
        "## Data-parallel neural network training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K0voYQTnSloK"
      },
      "source": [
        "### Preparing data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 166,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M3B-FAZOOKOH",
        "outputId": "1f815f28-b3f6-440d-d237-aa62c51e106c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-11 23:43:56.928815: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
            "2024-08-11 23:43:56.950628: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
            "2024-08-11 23:43:56.957397: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
            "2024-08-11 23:43:58.017597: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
          ]
        }
      ],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "data_dir = '/tmp/tfds'\n",
        "\n",
        "data, info = tfds.load(name=\"mnist\",\n",
        "                       data_dir=data_dir,\n",
        "                       as_supervised=True,\n",
        "                       with_info=True)\n",
        "\n",
        "data_train = data['train']\n",
        "data_test  = data['test']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 167,
      "metadata": {
        "id": "GQW9KLknOSc4"
      },
      "outputs": [],
      "source": [
        "HEIGHT = 28\n",
        "WIDTH  = 28\n",
        "CHANNELS = 1\n",
        "NUM_PIXELS = HEIGHT * WIDTH * CHANNELS\n",
        "NUM_LABELS = info.features['label'].num_classes\n",
        "NUM_DEVICES = jax.device_count()\n",
        "BATCH_SIZE  = 32"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 168,
      "metadata": {
        "id": "XEL-NZOjOV0_"
      },
      "outputs": [],
      "source": [
        "def preprocess(img, label):\n",
        "  \"\"\"Resize and preprocess images.\"\"\"\n",
        "  return (tf.cast(img, tf.float32)/255.0), label\n",
        "\n",
        "train_data = tfds.as_numpy(\n",
        "    data_train.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")\n",
        "test_data  = tfds.as_numpy(\n",
        "    data_test.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 169,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bcwCdU2-GU80",
        "outputId": "b303530c-c59a-4a87-ecb3-58ce29efda04"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "469"
            ]
          },
          "execution_count": 169,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ilU3LzAeSolR"
      },
      "source": [
        "### Preparing MLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 170,
      "metadata": {
        "id": "Q7B2PFjwOben"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jax import grad, jit, vmap, value_and_grad\n",
        "from jax import random\n",
        "from jax.nn import swish, logsumexp, one_hot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 171,
      "metadata": {
        "id": "aG1Ixe95OiIv"
      },
      "outputs": [],
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "PARAM_SCALE = 0.01"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 172,
      "metadata": {
        "id": "DIwXobuPOprS"
      },
      "outputs": [],
      "source": [
        "def init_network_params(sizes, key=random.PRNGKey(0), scale=1e-2):\n",
        "  \"\"\"Initialize all layers for a fully-connected neural network with given sizes\"\"\"\n",
        "\n",
        "  def random_layer_params(m, n, key, scale=1e-2):\n",
        "    \"\"\"A helper function to randomly initialize weights and biases of a dense layer\"\"\"\n",
        "    w_key, b_key = random.split(key)\n",
        "    return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n",
        "\n",
        "  keys = random.split(key, len(sizes))\n",
        "  return [random_layer_params(m, n, k, scale) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n",
        "\n",
        "init_params = init_network_params(LAYER_SIZES, random.PRNGKey(0), scale=PARAM_SCALE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 173,
      "metadata": {
        "id": "DjyD2_EmNlC5"
      },
      "outputs": [],
      "source": [
        "def predict(params, image):\n",
        "  \"\"\"Function for per-example predictions.\"\"\"\n",
        "  activations = image\n",
        "  for w, b in params[:-1]:\n",
        "    outputs = jnp.dot(w, activations) + b\n",
        "    activations = swish(outputs)\n",
        "\n",
        "  final_w, final_b = params[-1]\n",
        "  logits = jnp.dot(final_w, activations) + final_b\n",
        "  return logits\n",
        "\n",
        "batched_predict = vmap(predict, in_axes=(None, 0))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TdwS4P06Sr-P"
      },
      "source": [
        "### Loss and update functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 174,
      "metadata": {
        "id": "bwg667wGO3Yy"
      },
      "outputs": [],
      "source": [
        "INIT_LR = 1.0\n",
        "DECAY_RATE = 0.95\n",
        "DECAY_STEPS = 5\n",
        "NUM_EPOCHS  = 20"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 175,
      "metadata": {
        "id": "T7vWgrxUKqkH"
      },
      "outputs": [],
      "source": [
        "from functools import partial"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0V3nOttiO7da"
      },
      "outputs": [],
      "source": [
        "def loss(params, images, targets):\n",
        "  \"\"\"Categorical cross entropy loss function.\"\"\"\n",
        "  logits = batched_predict(params, images)\n",
        "  log_preds = logits - jnp.expand_dims(logsumexp(logits, axis=1), 1)  # logsumexp trick https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/\n",
        "  return -jnp.mean(targets*log_preds)\n",
        "\n",
        "@partial(jax.pmap, axis_name='devices', in_axes=(None, 0, 0, None), out_axes=(None,0))\n",
        "def update(params, x, y, epoch_number):\n",
        "  loss_value, grads = value_and_grad(loss)(params, x, y)\n",
        "  grads = [(jax.lax.psum(dw, 'devices'), jax.lax.psum(db, 'devices'))\n",
        "    for dw, db in grads]\n",
        "  lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS)\n",
        "  return [(w - lr * dw, b - lr * db)\n",
        "          for (w, b), (dw, db) in zip(params, grads)], loss_value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GKMrN2vOPLiW"
      },
      "source": [
        "### Section for debugging purposes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9l5wT9U0O-7J"
      },
      "outputs": [],
      "source": [
        "train_data_iter = iter(train_data)\n",
        "x, y = next(train_data_iter)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fm2cdhekRNFC",
        "outputId": "9f02cf6b-02ec-4db1-988a-499fd0073b10"
      },
      "outputs": [],
      "source": [
        "x.shape, y.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MpJ5tnZ4PQMN",
        "outputId": "86579337-bd23-4e49-a5c2-eeeaffc47810"
      },
      "outputs": [],
      "source": [
        "x = jnp.reshape(x, (NUM_DEVICES, BATCH_SIZE, NUM_PIXELS))\n",
        "y = jnp.reshape(one_hot(y, NUM_LABELS), (NUM_DEVICES, BATCH_SIZE, NUM_LABELS))\n",
        "x.shape, y.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vSKemd9_QW1_"
      },
      "outputs": [],
      "source": [
        "updated_params, loss_value = update(init_params, x, y, 0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E3VpMr--Rf67",
        "outputId": "71f14b13-5e80-40f3-f015-06834caa845a"
      },
      "outputs": [],
      "source": [
        "loss_value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4PFlq7EcS3q_"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zw4DIGl5InbS"
      },
      "outputs": [],
      "source": [
        "@jit\n",
        "def batch_accuracy(params, images, targets):\n",
        "  images = jnp.reshape(images, (len(images), NUM_PIXELS))\n",
        "  predicted_class = jnp.argmax(batched_predict(params, images), axis=1)\n",
        "  return jnp.mean(predicted_class == targets)\n",
        "\n",
        "def accuracy(params, data):\n",
        "  accs = []\n",
        "  for images, targets in data:\n",
        "    accs.append(batch_accuracy(params, images, targets))\n",
        "  return jnp.mean(jnp.array(accs))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "skYy98qTFUpP",
        "outputId": "3912fd05-dbf0-47e6-a46a-3650107a2192"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "params = init_params\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "  start_time = time.time()\n",
        "  losses = []\n",
        "  for x, y in train_data:\n",
        "    num_elements = len(y)\n",
        "    x = jnp.reshape(x, (NUM_DEVICES, num_elements//NUM_DEVICES, NUM_PIXELS))\n",
        "    y = jnp.reshape(one_hot(y, NUM_LABELS), (NUM_DEVICES, num_elements//NUM_DEVICES, NUM_LABELS))\n",
        "    params, loss_value = update(params, x, y, epoch)\n",
        "    losses.append(jnp.sum(loss_value))\n",
        "  epoch_time = time.time() - start_time\n",
        "\n",
        "  train_acc = accuracy(params, train_data)\n",
        "  test_acc = accuracy(params, test_data)\n",
        "  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
        "  print(\"Training set loss {}\".format(jnp.mean(jnp.array(losses))))\n",
        "  print(\"Training set accuracy {}\".format(train_acc))\n",
        "  print(\"Test set accuracy {}\".format(test_acc))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zacBQRSbNSkB"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}