Chapter-7/JAX_in_Action_Chapter_7_pmap.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## ⚖️ Choose A or B:"
      ],
      "metadata": {
        "id": "JCirjpcFSPan"
      }
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "N5lNqNEGtLyO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'"
      ],
      "metadata": {
        "id": "ypaH8OgftR_H"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "n8ip1VKgterO"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.devices(\"cpu\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tlh75iWutfVT",
        "outputId": "1e3f0edd-2820-445d-ce64-c095f889e2b7"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[CpuDevice(id=0),\n",
              " CpuDevice(id=1),\n",
              " CpuDevice(id=2),\n",
              " CpuDevice(id=3),\n",
              " CpuDevice(id=4),\n",
              " CpuDevice(id=5),\n",
              " CpuDevice(id=6),\n",
              " CpuDevice(id=7)]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## B: Setting up TPU"
      ],
      "metadata": {
        "id": "Rmk_qf6mKCZx"
      }
    },
    {
      "cell_type": "markdown",
      "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 !!"
      ],
      "metadata": {
        "id": "RMq70OC9wVHQ"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "YnSeGp15SDKi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# need for our later example with SPMD neural net training\n",
        "!pip install tensorflow"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FDq6YUpdSLNf",
        "outputId": "59875e4e-4fd8-4db6-a238-32ff0f8d33b9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\n",
            "Collecting tensorflow\n",
            "  Downloading tensorflow-2.13.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (479.6 MB)\n",
            "\u001b[K     |█████████████████████           | 315.1 MB 431 kB/s eta 0:06:22"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "IOPub data rate exceeded.\n",
            "The Jupyter server will temporarily stop sending output\n",
            "to the client in order to avoid crashing it.\n",
            "To change this limit, set the config variable\n",
            "`--ServerApp.iopub_data_rate_limit`.\n",
            "\n",
            "Current values:\n",
            "ServerApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
            "ServerApp.rate_limit_window=3.0 (secs)\n",
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[K     |████████████████████████████████| 479.6 MB 8.1 kB/s \n",
            "\u001b[?25hCollecting absl-py>=1.0.0\n",
            "  Downloading absl_py-2.0.0-py3-none-any.whl (130 kB)\n",
            "\u001b[K     |████████████████████████████████| 130 kB 100.0 MB/s \n",
            "\u001b[?25hCollecting astunparse>=1.6.0\n",
            "  Downloading astunparse-1.6.3-py2.py3-none-any.whl (12 kB)\n",
            "Collecting flatbuffers>=23.1.21\n",
            "  Downloading flatbuffers-23.5.26-py2.py3-none-any.whl (26 kB)\n",
            "Collecting gast<=0.4.0,>=0.2.1\n",
            "  Downloading gast-0.4.0-py3-none-any.whl (9.8 kB)\n",
            "Collecting google-pasta>=0.1.1\n",
            "  Downloading google_pasta-0.2.0-py3-none-any.whl (57 kB)\n",
            "\u001b[K     |████████████████████████████████| 57 kB 9.0 MB/s \n",
            "\u001b[?25hCollecting grpcio<2.0,>=1.24.3\n",
            "  Downloading grpcio-1.59.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.3 MB 84.0 MB/s \n",
            "\u001b[?25hCollecting h5py>=2.9.0\n",
            "  Downloading h5py-3.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.8 MB 93.5 MB/s \n",
            "\u001b[?25hCollecting keras<2.14,>=2.13.1\n",
            "  Downloading keras-2.13.1-py3-none-any.whl (1.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.7 MB 91.6 MB/s \n",
            "\u001b[?25hCollecting libclang>=13.0.0\n",
            "  Downloading libclang-16.0.6-py2.py3-none-manylinux2010_x86_64.whl (22.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 22.9 MB 95.3 MB/s \n",
            "\u001b[?25hCollecting numpy<=1.24.3,>=1.22\n",
            "  Downloading numpy-1.24.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 17.3 MB 93.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: opt-einsum>=2.3.2 in ./.local/lib/python3.8/site-packages (from tensorflow) (3.3.0)\n",
            "Requirement already satisfied: packaging in /usr/lib/python3/dist-packages (from tensorflow) (20.3)\n",
            "Collecting protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3\n",
            "  Downloading protobuf-4.24.4-cp37-abi3-manylinux2014_x86_64.whl (311 kB)\n",
            "\u001b[K     |████████████████████████████████| 311 kB 84.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tensorflow) (62.3.2)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/lib/python3/dist-packages (from tensorflow) (1.14.0)\n",
            "Collecting tensorboard<2.14,>=2.13\n",
            "  Downloading tensorboard-2.13.0-py3-none-any.whl (5.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.6 MB 94.8 MB/s \n",
            "\u001b[?25hCollecting tensorflow-estimator<2.14,>=2.13.0\n",
            "  Downloading tensorflow_estimator-2.13.0-py2.py3-none-any.whl (440 kB)\n",
            "\u001b[K     |████████████████████████████████| 440 kB 101.9 MB/s \n",
            "\u001b[?25hCollecting termcolor>=1.1.0\n",
            "  Downloading termcolor-2.3.0-py3-none-any.whl (6.9 kB)\n",
            "Collecting typing-extensions<4.6.0,>=3.6.6\n",
            "  Downloading typing_extensions-4.5.0-py3-none-any.whl (27 kB)\n",
            "Collecting wrapt>=1.11.0\n",
            "  Downloading wrapt-1.15.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (81 kB)\n",
            "\u001b[K     |████████████████████████████████| 81 kB 13.4 MB/s \n",
            "\u001b[?25hCollecting tensorflow-io-gcs-filesystem>=0.23.1; platform_machine != \"arm64\" or platform_system != \"Darwin\"\n",
            "  Downloading tensorflow_io_gcs_filesystem-0.34.0-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (2.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.4 MB 97.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wheel<1.0,>=0.23.0 in /usr/lib/python3/dist-packages (from astunparse>=1.6.0->tensorflow) (0.34.2)\n",
            "Collecting google-auth<3,>=1.6.3\n",
            "  Downloading google_auth-2.23.3-py2.py3-none-any.whl (182 kB)\n",
            "\u001b[K     |████████████████████████████████| 182 kB 108.3 MB/s \n",
            "\u001b[?25hCollecting google-auth-oauthlib<1.1,>=0.5\n",
            "  Downloading google_auth_oauthlib-1.0.0-py2.py3-none-any.whl (18 kB)\n",
            "Collecting markdown>=2.6.8\n",
            "  Downloading Markdown-3.5-py3-none-any.whl (101 kB)\n",
            "\u001b[K     |████████████████████████████████| 101 kB 17.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests<3,>=2.21.0 in ./.local/lib/python3.8/site-packages (from tensorboard<2.14,>=2.13->tensorflow) (2.31.0)\n",
            "Collecting tensorboard-data-server<0.8.0,>=0.7.0\n",
            "  Downloading tensorboard_data_server-0.7.2-py3-none-any.whl (2.4 kB)\n",
            "Collecting werkzeug>=1.0.1\n",
            "  Downloading werkzeug-3.0.1-py3-none-any.whl (226 kB)\n",
            "\u001b[K     |████████████████████████████████| 226 kB 98.9 MB/s \n",
            "\u001b[?25hCollecting cachetools<6.0,>=2.0.0\n",
            "  Downloading cachetools-5.3.2-py3-none-any.whl (9.3 kB)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/lib/python3/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.14,>=2.13->tensorflow) (0.2.1)\n",
            "Collecting rsa<5,>=3.1.4\n",
            "  Downloading rsa-4.9-py3-none-any.whl (34 kB)\n",
            "Collecting requests-oauthlib>=0.7.0\n",
            "  Downloading requests_oauthlib-1.3.1-py2.py3-none-any.whl (23 kB)\n",
            "Requirement already satisfied: importlib-metadata>=4.4; python_version < \"3.10\" in ./.local/lib/python3.8/site-packages (from markdown>=2.6.8->tensorboard<2.14,>=2.13->tensorflow) (6.8.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (1.25.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2019.11.28)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in ./.local/lib/python3.8/site-packages (from werkzeug>=1.0.1->tensorboard<2.14,>=2.13->tensorflow) (2.1.3)\n",
            "Requirement already satisfied: pyasn1>=0.1.3 in /usr/lib/python3/dist-packages (from rsa<5,>=3.1.4->google-auth<3,>=1.6.3->tensorboard<2.14,>=2.13->tensorflow) (0.4.2)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<1.1,>=0.5->tensorboard<2.14,>=2.13->tensorflow) (3.1.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/lib/python3/dist-packages (from importlib-metadata>=4.4; python_version < \"3.10\"->markdown>=2.6.8->tensorboard<2.14,>=2.13->tensorflow) (1.0.0)\n",
            "Installing collected packages: absl-py, astunparse, flatbuffers, gast, google-pasta, grpcio, numpy, h5py, keras, libclang, protobuf, cachetools, rsa, google-auth, requests-oauthlib, google-auth-oauthlib, markdown, tensorboard-data-server, werkzeug, tensorboard, tensorflow-estimator, termcolor, typing-extensions, wrapt, tensorflow-io-gcs-filesystem, tensorflow\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 1.24.4\n",
            "    Uninstalling numpy-1.24.4:\n",
            "      Successfully uninstalled numpy-1.24.4\n",
            "  Attempting uninstall: typing-extensions\n",
            "    Found existing installation: typing-extensions 4.8.0\n",
            "    Uninstalling typing-extensions-4.8.0:\n",
            "      Successfully uninstalled typing-extensions-4.8.0\n",
            "Successfully installed absl-py-2.0.0 astunparse-1.6.3 cachetools-5.3.2 flatbuffers-23.5.26 gast-0.4.0 google-auth-2.23.3 google-auth-oauthlib-1.0.0 google-pasta-0.2.0 grpcio-1.59.2 h5py-3.10.0 keras-2.13.1 libclang-16.0.6 markdown-3.5 numpy-1.24.3 protobuf-4.24.4 requests-oauthlib-1.3.1 rsa-4.9 tensorboard-2.13.0 tensorboard-data-server-0.7.2 tensorflow-2.13.1 tensorflow-estimator-2.13.0 tensorflow-io-gcs-filesystem-0.34.0 termcolor-2.3.0 typing-extensions-4.5.0 werkzeug-3.0.1 wrapt-1.15.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow_datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VMCFxcZcS0u0",
        "outputId": "d042dd40-41ec-4ccc-dfb9-dd913035d37a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\n",
            "Collecting tensorflow_datasets\n",
            "  Downloading tensorflow_datasets-4.9.2-py3-none-any.whl (5.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.4 MB 4.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: absl-py in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.0.0)\n",
            "Collecting array-record\n",
            "  Downloading array_record-0.4.0-py38-none-any.whl (3.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.0 MB 93.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: click in /usr/lib/python3/dist-packages (from tensorflow_datasets) (7.0)\n",
            "Collecting dm-tree\n",
            "  Downloading dm_tree-0.1.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (152 kB)\n",
            "\u001b[K     |████████████████████████████████| 152 kB 105.7 MB/s \n",
            "\u001b[?25hCollecting etils[enp,epath]>=0.9.0\n",
            "  Downloading etils-1.3.0-py3-none-any.whl (126 kB)\n",
            "\u001b[K     |████████████████████████████████| 126 kB 110.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (1.24.3)\n",
            "Collecting promise\n",
            "  Downloading promise-2.3.tar.gz (19 kB)\n",
            "Requirement already satisfied: protobuf>=3.20 in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (4.24.4)\n",
            "Requirement already satisfied: psutil in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (5.9.6)\n",
            "Requirement already satisfied: requests>=2.19.0 in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.31.0)\n",
            "Collecting tensorflow-metadata\n",
            "  Downloading tensorflow_metadata-1.14.0-py3-none-any.whl (28 kB)\n",
            "Requirement already satisfied: termcolor in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.3.0)\n",
            "Collecting toml\n",
            "  Downloading toml-0.10.2-py2.py3-none-any.whl (16 kB)\n",
            "Collecting tqdm\n",
            "  Downloading tqdm-4.66.1-py3-none-any.whl (78 kB)\n",
            "\u001b[K     |████████████████████████████████| 78 kB 9.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wrapt in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (1.15.0)\n",
            "Requirement already satisfied: importlib-resources; python_version < \"3.9\" in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (6.1.0)\n",
            "Requirement already satisfied: typing_extensions; extra == \"epath\" in ./.local/lib/python3.8/site-packages (from etils[enp,epath]>=0.9.0->tensorflow_datasets) (4.5.0)\n",
            "Requirement already satisfied: zipp; extra == \"epath\" in /usr/lib/python3/dist-packages (from etils[enp,epath]>=0.9.0->tensorflow_datasets) (1.0.0)\n",
            "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from promise->tensorflow_datasets) (1.14.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (1.25.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2019.11.28)\n",
            "Collecting googleapis-common-protos<2,>=1.52.0\n",
            "  Downloading googleapis_common_protos-1.61.0-py2.py3-none-any.whl (230 kB)\n",
            "\u001b[K     |████████████████████████████████| 230 kB 104.1 MB/s \n",
            "\u001b[?25hBuilding wheels for collected packages: promise\n",
            "  Building wheel for promise (setup.py) ... \u001b[?25l-\b \b\\\b \b|\b \b/\b \b-\b \b\\\b \bdone\n",
            "\u001b[?25h  Created wheel for promise: filename=promise-2.3-py3-none-any.whl size=21485 sha256=8a939e3a2cf29d11f7857c4bd0188887243877e6695db8094476a6d6ab6028b9\n",
            "  Stored in directory: /home/grigo/.cache/pip/wheels/54/aa/01/724885182f93150035a2a91bce34a12877e8067a97baaf5dc8\n",
            "Successfully built promise\n",
            "\u001b[31mERROR: tensorflow-metadata 1.14.0 has requirement absl-py<2.0.0,>=0.9, but you'll have absl-py 2.0.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow-metadata 1.14.0 has requirement protobuf<4.21,>=3.20.3, but you'll have protobuf 4.24.4 which is incompatible.\u001b[0m\n",
            "Installing collected packages: etils, array-record, dm-tree, promise, googleapis-common-protos, tensorflow-metadata, toml, tqdm, tensorflow-datasets\n",
            "Successfully installed array-record-0.4.0 dm-tree-0.1.8 etils-1.3.0 googleapis-common-protos-1.61.0 promise-2.3 tensorflow-datasets-4.9.2 tensorflow-metadata-1.14.0 toml-0.10.2 tqdm-4.66.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.lib import xla_bridge\n",
        "print(xla_bridge.get_backend().platform)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2aBy_QGmu-61",
        "outputId": "42573514-97a2-43b5-bd84-f29ac4a15cb8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "7c2IXWjRvFbj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.local_devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_qLo4vNRvHc-",
        "outputId": "f1a64216-4f71-4dc0-b052-6405b63aab8a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
              " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
              " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
              " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Parallelizing a function"
      ],
      "metadata": {
        "id": "wzZlvWsjKIC8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "qxtvsXER1Ums"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dot(jnp.array([1., 1., 1.]), jnp.array([1., 2., -1]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UuEi71y-1iDg",
        "outputId": "2d1a4f77-51da-4e3c-ea3a-b39c069ec2c8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(2., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generate some **large** random arrays"
      ],
      "metadata": {
        "id": "f6uqvbVsKLdp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import random"
      ],
      "metadata": {
        "id": "4cVOwuJyv6yO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rng_key = random.PRNGKey(42)"
      ],
      "metadata": {
        "id": "EUTNwutCwiUl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vs = random.normal(rng_key, shape=(20_000_000,3))"
      ],
      "metadata": {
        "id": "MSGxkzRGwmiy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s = vs[:10_000_000,:]\n",
        "v2s = vs[10_000_000:,:]"
      ],
      "metadata": {
        "id": "zLYOfhF19gI4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.shape, v2s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oBMFv_cc-GV6",
        "outputId": "a6530e75-e9e9-44b8-cebb-f0a0b82ed8db"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((10000000, 3), (10000000, 3))"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make a compiled vectorized function as a baseline:"
      ],
      "metadata": {
        "id": "fHT8Pi7YPNvC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dot_batched = jax.jit(jax.vmap(dot)).lower(v1s,v2s).compile() # you can use AOT-compilation"
      ],
      "metadata": {
        "id": "KsRK7fsv92RN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit x_vmap = dot_batched(v1s, v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b-NUKyM6gO1R",
        "outputId": "a82076bc-4a26-4e92-c827-7cd762f5eda7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.91 ms ± 52.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_batched = jax.jit(jax.vmap(dot)) # or JIT-compilation"
      ],
      "metadata": {
        "id": "yPKa12rgEan4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_vmap = dot_batched(v1s, v2s)"
      ],
      "metadata": {
        "id": "tQd2syYJllpQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_vmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9uZUaB6nE5pT",
        "outputId": "6a81e3a1-ca18-4543-e786-f6eb13a5d1cb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_parallel = jax.pmap(dot)"
      ],
      "metadata": {
        "id": "ZkmFPNA-hjbR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = dot_parallel(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 488
        },
        "id": "J4s0rZIVh0RC",
        "outputId": "6301885b-5859-45ac-ac53-d76a18ecb813"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[17], 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~/.local/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py:993\u001b[0m, in \u001b[0;36mUnloadedPmapExecutable.from_hlo\u001b[0;34m(hlo, pci, replicas, shards, tuple_args, unordered_effects, ordered_effects, host_callbacks, keepalive, jaxpr_debug_info, compiler_options)\u001b[0m\n\u001b[1;32m    990\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    991\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    992\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--> 993\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    994\u001b[0m                               xb\u001b[38;5;241m.\u001b[39mdevice_count(pci\u001b[38;5;241m.\u001b[39mbackend),\n\u001b[1;32m    995\u001b[0m                               replicas\u001b[38;5;241m.\u001b[39mnum_global_replicas))\n\u001b[1;32m    996\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    997\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    998\u001b[0m \u001b[38;5;66;03m# match jax.devices().\u001b[39;00m\n\u001b[1;32m    999\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   1000\u001b[0m \u001b[38;5;66;03m# each host is responsible for a contiguous set of replicas.\u001b[39;00m\n\u001b[1;32m   1001\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   1002\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   1003\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 8 XLA devices are available (num_replicas=10000000)"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1DCjJciai7Z_",
        "outputId": "be77bca5-7898-4418-838b-e0dc336fd0d3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = v1s.reshape((8, v1s.shape[0]//8, v1s.shape[1]))\n",
        "v2sp = v2s.reshape((8, v2s.shape[0]//8, v2s.shape[1]))    # we use integer division here"
      ],
      "metadata": {
        "id": "0zaUwp4tjroC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ntlAFzYAi9_s",
        "outputId": "04ae1775-877d-4746-f5a4-8a0cd503ec89"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1s[0,:], v1sp[0,0,:]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ttr1LezXIvos",
        "outputId": "56152362-0aae-4dde-b8b7-40943627c53e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Array([1.456546  , 1.1449531 , 0.02485494], dtype=float32),\n",
              " Array([1.456546  , 1.1449531 , 0.02485494], dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1s[1250000,:], v1sp[1,0,:]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xQJAXG1MI6ja",
        "outputId": "e26e2375-5d53-4bf3-c395-3b6f7cddd297"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Array([ 0.50284576,  0.38336843, -0.35499337], dtype=float32),\n",
              " Array([ 0.50284576,  0.38336843, -0.35499337], dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "kED5PhIvj_Yz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C4SjsR23KaYm",
        "outputId": "b8396ca5-4be1-449c-ecea-ba5288bc266d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8,)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_parallel = jax.pmap(jax.vmap(dot))"
      ],
      "metadata": {
        "id": "-5sG0zdjkNwR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "gCI71XoukYOv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W4oj2g-bkgbM",
        "outputId": "e437750a-d162-4749-af79-4ab67c98bc1f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "type(x_pmap)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dyIOI3cakjL0",
        "outputId": "0e0cebc0-8e86-4961-fbe0-9ef3ccc50936"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "jaxlib.xla_extension.ArrayImpl"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1W2vIrDxklDX",
        "outputId": "cff17983-2c42-4683-a47a-346921b3392c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.all(x_pmap == x_vmap)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "38UCSdMLmPYS",
        "outputId": "00a13cde-c89c-4206-dc20-ec650e88f086"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit xp = dot_parallel(v1sp,v2sp).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k2jjZe-mo47n",
        "outputId": "76ab37a2-2035-4a8d-e981-419ead1dcb10"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "9.9 ms ± 49.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-_jLRzyBRDTJ",
        "outputId": "9d5e09f6-e779-475d-be17-b812d0aab908"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 1.1981742 , -3.0097814 , -2.539115  , ..., -2.7281013 ,\n",
              "        0.26364914,  1.9583547 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_vmap"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TxECifZoRGYY",
        "outputId": "ca2afae7-9e4a-4238-9cb8-14e669c2e680"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 1.1981742 , -3.0097814 , -2.539115  , ..., -2.7281013 ,\n",
              "        0.26364914,  1.9583547 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A very simple example when you can easily switch vmap and pmap:"
      ],
      "metadata": {
        "id": "ReSRcPnlR_7g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "R28QzddqW6Xm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vs = random.normal(rng_key, shape=(16,3))\n",
        "v1s = vs[:8,:]\n",
        "v2s = vs[8:,:]"
      ],
      "metadata": {
        "id": "giEtn52mSIk-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.vmap(dot)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LfT1rhKJTgwq",
        "outputId": "63bb34b2-4461-4bfe-d853-bea782581fbf"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.pmap(dot)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z46RrWmhTnEK",
        "outputId": "5972ad49-cb7b-4d10-9a2c-7d453c675035"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_v = jax.jit(jax.vmap(dot))\n",
        "x = dot_v(v1s,v2s)"
      ],
      "metadata": {
        "id": "vNVXUxJpXLaB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pjo = jax.jit(jax.pmap(dot))\n",
        "x = dot_pjo(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TbZv85rgXRR_",
        "outputId": "0bd2420a-91d6-4299-844e-7d893344c667"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/home/grigo/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:290: 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"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pji = jax.pmap(jax.jit(dot))\n",
        "x = dot_pji(v1s,v2s)"
      ],
      "metadata": {
        "id": "OHXNHJTYYJb5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dot_p = jax.pmap(dot)\n",
        "x = dot_p(v1s,v2s)"
      ],
      "metadata": {
        "id": "KNrEftwLXZlw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit dot_v(v1s,v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pm9Q-E4UWcHs",
        "outputId": "56228848-94bf-4e56-bf3e-eb024ad4ccb0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "122 µs ± 5.53 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit dot_pjo(v1s,v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IJtuPjyBWiBp",
        "outputId": "c1e1e811-cb0f-4624-cb1e-0400270f6c50"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.08 ms ± 30.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit dot_pji(v1s,v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q6q8J8TCYQho",
        "outputId": "a3be9b08-b156-44c9-afb0-793a5789a711"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.51 ms ± 43 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit dot_p(v1s,v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KWDwONE-Xsjn",
        "outputId": "13ff8763-b842-464a-d4e6-85f1a6a87e9e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.54 ms ± 51 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(dot_v)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lJFuKu1XAhFj",
        "outputId": "7c2229a5-0767-44af-89b1-7aa8f6f7d0e0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[8,3]\u001b[39m b\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[8]\u001b[39m = pjit[\n",
              "      jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; d\u001b[35m:f32[8,3]\u001b[39m e\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "          \u001b[39m\u001b[22m\u001b[22mf\u001b[35m:f32[8]\u001b[39m = dot_general[dimension_numbers=(([1], [1]), ([0], [0]))] d e\n",
              "        \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(f,) }\n",
              "      name=dot\n",
              "    ] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 48
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(dot_pjo)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PJE-4HWRAnko",
        "outputId": "c60f2960-ac34-4733-f137-5f53f57da97a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[8,3]\u001b[39m b\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[8]\u001b[39m = pjit[\n",
              "      jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; d\u001b[35m:f32[8,3]\u001b[39m e\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "          \u001b[39m\u001b[22m\u001b[22mf\u001b[35m:f32[8]\u001b[39m = xla_pmap[\n",
              "            axis_name=<axis 0x7f42a91d2b80>\n",
              "            axis_size=8\n",
              "            backend=None\n",
              "            call_jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; g\u001b[35m:f32[3]\u001b[39m h\u001b[35m:f32[3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "                \u001b[39m\u001b[22m\u001b[22mi\u001b[35m:f32[]\u001b[39m = dot_general[dimension_numbers=(([0], [0]), ([], []))] g\n",
              "                  h\n",
              "              \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(i,) }\n",
              "            devices=None\n",
              "            donated_invars=(False, False)\n",
              "            global_axis_size=8\n",
              "            in_axes=(0, 0)\n",
              "            is_explicit_global_axis_size=False\n",
              "            name=dot\n",
              "            out_axes=(0,)\n",
              "          ] d e\n",
              "        \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(f,) }\n",
              "      name=dot\n",
              "    ] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(dot_pji)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VH60u2geApK3",
        "outputId": "6881a604-42cb-44bf-b27e-e52d8c87f1ca"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[8,3]\u001b[39m b\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[8]\u001b[39m = xla_pmap[\n",
              "      axis_name=<axis 0x7f42a8f09900>\n",
              "      axis_size=8\n",
              "      backend=None\n",
              "      call_jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; d\u001b[35m:f32[3]\u001b[39m e\u001b[35m:f32[3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "          \u001b[39m\u001b[22m\u001b[22mf\u001b[35m:f32[]\u001b[39m = pjit[\n",
              "            jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; g\u001b[35m:f32[3]\u001b[39m h\u001b[35m:f32[3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "                \u001b[39m\u001b[22m\u001b[22mi\u001b[35m:f32[]\u001b[39m = dot_general[dimension_numbers=(([0], [0]), ([], []))] g\n",
              "                  h\n",
              "              \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(i,) }\n",
              "            name=dot\n",
              "          ] d e\n",
              "        \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(f,) }\n",
              "      devices=None\n",
              "      donated_invars=(False, False)\n",
              "      global_axis_size=8\n",
              "      in_axes=(0, 0)\n",
              "      is_explicit_global_axis_size=False\n",
              "      name=dot\n",
              "      out_axes=(0,)\n",
              "    ] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 50
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(dot_p)(v1s,v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Alc3_10BAotr",
        "outputId": "7561b326-de10-4f1a-a264-06085045a764"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[8,3]\u001b[39m b\u001b[35m:f32[8,3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[8]\u001b[39m = xla_pmap[\n",
              "      axis_name=<axis 0x7f42a91d2b80>\n",
              "      axis_size=8\n",
              "      backend=None\n",
              "      call_jaxpr={ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; d\u001b[35m:f32[3]\u001b[39m e\u001b[35m:f32[3]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "          \u001b[39m\u001b[22m\u001b[22mf\u001b[35m:f32[]\u001b[39m = dot_general[dimension_numbers=(([0], [0]), ([], []))] d e\n",
              "        \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(f,) }\n",
              "      devices=None\n",
              "      donated_invars=(False, False)\n",
              "      global_axis_size=8\n",
              "      in_axes=(0, 0)\n",
              "      is_explicit_global_axis_size=False\n",
              "      name=dot\n",
              "      out_axes=(0,)\n",
              "    ] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 51
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Controlling pmap() behavior"
      ],
      "metadata": {
        "id": "rB5gg9rvePQb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Using in_axes parameter"
      ],
      "metadata": {
        "id": "1_dx8YVoGb-H"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Using a small array just for demonstration purposes"
      ],
      "metadata": {
        "id": "82lZ7bqNT4TN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vs = random.normal(rng_key, shape=(16,3))\n",
        "v1s = vs[:8,:]\n",
        "v2s = vs[8:,:]"
      ],
      "metadata": {
        "id": "WKUM9E0jT3sx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "pmydsax3UEkA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "A default value:"
      ],
      "metadata": {
        "id": "YxGyiGP5mlcK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pmapped = jax.pmap(dot, in_axes=(0,0))"
      ],
      "metadata": {
        "id": "_anvIZEhj-bk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pmapped(v1s, v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZjEzc-wedMDq",
        "outputId": "5791ec39-40b3-4839-ef9e-9b5c056fd850"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 55
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "What if one of our arrays is transposed (the mapping axis is the second one)?"
      ],
      "metadata": {
        "id": "EdvjNEkoUN-l"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.T.shape, v2s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3CNhCLoZal3p",
        "outputId": "e08db0a7-c13a-4c4a-bb1a-cef61c2c75dc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 8), (8, 3))"
            ]
          },
          "metadata": {},
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.pmap(dot, in_axes=(1,0))(v1s.T, v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jr-GZ4zmUNXa",
        "outputId": "21f6a200-9d6f-4d13-acfb-e0f1a40a90fd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pmapped(v1s.T, v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 332
        },
        "id": "4ZTIPUEyUeqJ",
        "outputId": "7a0b1619-820c-4dd9-dae7-07e4fa0a44d5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[58], 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~/.local/lib/python3.8/site-packages/jax/_src/api.py:1317\u001b[0m, in \u001b[0;36m_mapped_axis_size\u001b[0;34m(fn, tree, vals, dims, name)\u001b[0m\n\u001b[1;32m   1315\u001b[0m   \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1316\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-> 1317\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,8];\n  * one axis had size 8: axis 0 of argument v2 of type float32[8,3]"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.T.shape, v2s.T.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QQub20ZTauM9",
        "outputId": "729a6a83-f3e3-42a6-8254-13005140337d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 8), (3, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 59
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.pmap(dot, in_axes=(1,1))(v1s.T, v2s.T)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HXY3CZ9uUjIb",
        "outputId": "fb81fc11-cdda-46a1-f31b-1161cbf3a51f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 60
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_pmapped(v1s.T, v2s.T)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rgsH99-1UoWA",
        "outputId": "5fb5baa3-d7cf-4680-b264-307542483174"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.00782511, -0.31532776,  1.6295187 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 61
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A more complicated case:"
      ],
      "metadata": {
        "id": "gqz7L6uSmqPc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def scaled_dot(v1, v2, koeff):\n",
        "  return koeff*jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "q7ZB2g1GdWJz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s_ = v1s\n",
        "v2s_ = v2s.T\n",
        "k = 1.0"
      ],
      "metadata": {
        "id": "yjHKGd1-diSw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s_.shape, v2s_.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LU0SfSJieggS",
        "outputId": "43d5171f-8a50-4e0e-f2c2-aa61dcd9d1eb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 3), (3, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 64
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot)"
      ],
      "metadata": {
        "id": "ymOex9oHemQ8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Default values do not work:"
      ],
      "metadata": {
        "id": "Q-pqO8b_muy4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped(v1s_, v2s_, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 523
        },
        "id": "dbFc89nYfaH8",
        "outputId": "22995d8e-1604-43b3-d6df-f2851afeccd0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/api.py:1260\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   1259\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1260\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   1261\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[66], 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~/.local/lib/python3.8/site-packages/jax/_src/api.py:1264\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   1262\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   1263\u001b[0m \u001b[38;5;66;03m# TODO(mattjj): better error message here\u001b[39;00m\n\u001b[0;32m-> 1264\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   1265\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   1266\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   1267\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 ())"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot, in_axes=(0,1,None))"
      ],
      "metadata": {
        "id": "9dTFcAHQfdVf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped(v1s_, v2s_, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A2f452A9flEO",
        "outputId": "0522f2ee-a1b2-4940-d05d-1a5a34cb06bb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 68
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Using more complex parameter structure:"
      ],
      "metadata": {
        "id": "PVSBhIYAmyi4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def scaled_dot(data, koeff):\n",
        "  return koeff*jnp.vdot(data['a'], data['b'])"
      ],
      "metadata": {
        "id": "6cuNZDUZiMEH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped = jax.pmap(scaled_dot, in_axes=({'a':0,'b':1},None))"
      ],
      "metadata": {
        "id": "pfdMJ7eeijRq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scaled_dot_pmapped({'a':v1s_, 'b': v2s_}, k)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lMjD-rsCivBv",
        "outputId": "1441d7fc-62c7-4223-c49f-1c09708f8362"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.51048726, -0.7174605 , -0.20105815, -0.26437205, -1.3696793 ,\n",
              "        2.744793  ,  1.7936493 , -1.1743435 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 71
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Using out_axes parameter"
      ],
      "metadata": {
        "id": "72VWPGXrm5du"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def scale(v, koeff):\n",
        "  return koeff*v"
      ],
      "metadata": {
        "id": "dv3wsrlJm7BB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scale_pmapped = jax.pmap(scale,\n",
        "                         in_axes=(0,None),\n",
        "                         out_axes=(1))"
      ],
      "metadata": {
        "id": "C34aK9bgnBed"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = scale_pmapped(v1s, 2.0)"
      ],
      "metadata": {
        "id": "DodCkW4oiz2p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.shape, res.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0UTVvz2ZgPQj",
        "outputId": "79ac19dc-00cb-4121-c137-dcfe0f00a476"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 3), (3, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 75
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "scale_pmapped = jax.pmap(scale, in_axes=(0,None))"
      ],
      "metadata": {
        "id": "DGrRn25qnYrj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scale_pmapped(v1s, 2.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4qzPVn_znZ3o",
        "outputId": "95220e2d-757a-4bf3-a5e9-048dfbd04d74"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[-0.44077486, -1.7589345 , -1.7889494 ],\n",
              "       [-0.5077695 , -0.6062154 ,  0.8972974 ],\n",
              "       [-1.2127566 ,  0.33352432, -0.9126665 ],\n",
              "       [-0.41174212,  1.8439273 ,  1.8103753 ],\n",
              "       [ 0.550015  , -2.7737749 , -1.8786148 ],\n",
              "       [-0.13278721, -2.1072452 ,  2.6899092 ],\n",
              "       [-0.3347797 ,  2.1231527 , -3.6858454 ],\n",
              "       [-0.24115095,  2.8006103 , -1.3478842 ]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "scale_pmapped = jax.pmap(scale,\n",
        "                         in_axes=(0,None),\n",
        "                         out_axes=None)"
      ],
      "metadata": {
        "id": "K5Jps90J6HLw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A more real-life case with larger tensors and mixing vmap and pmap"
      ],
      "metadata": {
        "id": "ru8LmRBtYavl"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "wGaZtjJ1hqw7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1s.shape, v2s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hasxnv11h-bb",
        "outputId": "85d5d103-ddcc-48b4-9e42-ba1d04254cba"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 10000000), (3, 10000000))"
            ]
          },
          "metadata": {},
          "execution_count": 80
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = v1s.reshape((v1s.shape[0], 8, v1s.shape[1]//8))\n",
        "v2sp = v2s.reshape((v2s.shape[0], 8, v2s.shape[1]//8))\n",
        "v1sp.shape, v2sp.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lOkDY2wromwx",
        "outputId": "ad86e702-3b3d-4131-8972-3391a08181d1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 8, 1250000), (3, 8, 1250000))"
            ]
          },
          "metadata": {},
          "execution_count": 81
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_parallel = jax.pmap(\n",
        "    jax.vmap(dot, in_axes=(1,1)),\n",
        "    in_axes=(1,1)\n",
        ")"
      ],
      "metadata": {
        "id": "wRgFXKt-o0aQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "tEa0RhW0pW5X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wxz83U68pb0I",
        "outputId": "bd89f3c7-c246-4d24-cada-cb9fbae433a5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000)"
            ]
          },
          "metadata": {},
          "execution_count": 84
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1kEHvzk5pmLe",
        "outputId": "8f200aed-ae8e-4e8d-d1cb-57112a5dc17d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 85
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.all(x_pmap == x_vmap)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m-AaCUnGpp4l",
        "outputId": "d89f63b7-2d1b-4514-ccb2-58822b37da22"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 86
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap[:3]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yd_Y1-YOpvdU",
        "outputId": "6d9b62da-9cfe-4533-8792-1188bb0a778d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 1.1981742, -3.0097814, -2.539115 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_vmap[:3]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0VbvxplMpy22",
        "outputId": "ae25a93b-7279-4f71-c983-8c2954637c3b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 1.1981742, -3.0097814, -2.539115 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 88
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using collective ops and the axis_name parameter"
      ],
      "metadata": {
        "id": "7zxdlwUP14or"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Normalization examples"
      ],
      "metadata": {
        "id": "dF2ONW2AhMVd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "A small array example"
      ],
      "metadata": {
        "id": "POmEA_DwhRD6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(8))"
      ],
      "metadata": {
        "id": "DrWA7KPZ196e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "arr"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iGDJGEfn245M",
        "outputId": "df778b2d-75b2-4f49-b829-20c8bc18c460"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 90
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "norm = jax.pmap(\n",
        "    lambda x: x/jax.lax.psum(x, axis_name='p'),\n",
        "    axis_name='p')"
      ],
      "metadata": {
        "id": "AYLt3IU8_pFE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "norm(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2KXDuNJi_zv7",
        "outputId": "fe8f62ce-4762-4f11-f5ca-2f0c0b3627d5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0.        , 0.03571429, 0.07142857, 0.10714287, 0.14285715,\n",
              "       0.17857143, 0.21428573, 0.25      ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 93
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(norm(arr))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5VzX8Gc4_1vo",
        "outputId": "1f172ea5-875d-4e0e-9a97-798311962101"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 94
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "A large array (more elements than XLA devices) example"
      ],
      "metadata": {
        "id": "BAsivlcghTHm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(200))"
      ],
      "metadata": {
        "id": "MEDDfwp21SH8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "arr = arr.reshape(8, 25)\n",
        "arr.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4otHYX5R1VKR",
        "outputId": "392dd228-748e-436c-de59-9c2bfe1f9232"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 25)"
            ]
          },
          "metadata": {},
          "execution_count": 96
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "norm = jax.pmap(\n",
        "    lambda x: x/jax.lax.psum(jnp.sum(x), axis_name='p'),\n",
        "    axis_name='p')"
      ],
      "metadata": {
        "id": "4n21Oyvp1Z-t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fCHy2q-n1eah",
        "outputId": "678f2166-ad1b-4bbc-a01e-bed5a1e9547c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 25)"
            ]
          },
          "metadata": {},
          "execution_count": 98
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(narr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SUEfvESE4bn_",
        "outputId": "8bf0afd3-678b-4a54-8664-38e5f782881e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Using groups"
      ],
      "metadata": {
        "id": "0lFI5h1XhyqI"
      }
    },
    {
      "cell_type": "code",
      "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], [4,5], [6,7]]\n",
        "    ),\n",
        "    axis_name='p')"
      ],
      "metadata": {
        "id": "y3Zt5UpD-l3L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "58qiMx2e_E8l",
        "outputId": "730d5df7-69f4-4548-b16f-5d54178d7673"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 25)"
            ]
          },
          "metadata": {},
          "execution_count": 101
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(narr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iDxWZ8xN_VqZ",
        "outputId": "01ad325a-31e1-4621-92a0-d7e205030567"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(4., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 102
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(narr[:2]), jnp.sum(narr[2:4]), jnp.sum(narr[4:6]), jnp.sum(narr[6:])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zTRiyInw_lmT",
        "outputId": "600d2d59-5c86-49cb-ba1f-8beaff122cc9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Array(1., dtype=float32),\n",
              " Array(1., dtype=float32),\n",
              " Array(1.0000001, dtype=float32),\n",
              " Array(1., dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 103
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Using nested pmap() and vmap()"
      ],
      "metadata": {
        "id": "mAFZXq76DQbe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(200))\n",
        "arr = arr.reshape(8, 25)\n",
        "arr.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xY7giT0-DP_C",
        "outputId": "15c11967-83aa-472b-97b7-e7627735b582"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 25)"
            ]
          },
          "metadata": {},
          "execution_count": 104
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Understanding different axes"
      ],
      "metadata": {
        "id": "uNs_Sri4hf7c"
      }
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "S54zj0qMFxNb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "f(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZiMDUbLlFzl8",
        "outputId": "ba96c985-7a70-4685-d9e0-228d010e25e1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199],\n",
              "       [175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187,\n",
              "        188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199]],      dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 106
        }
      ]
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "47d7KE6PfwqW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "f(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-blZXHawf0sp",
        "outputId": "822ebd77-c216-42b1-b6f9-d404bbbe4526"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "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],\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],\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],\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],\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],\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],\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],\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]],      dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 108
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Using both axes"
      ],
      "metadata": {
        "id": "DfLX3pMphlg1"
      }
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "WhVL-R8qDd7_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "f(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rpNlgksOEEpJ",
        "outputId": "214388ee-096c-4b3a-ca45-f5e61f0f1b4a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0.13714285, 0.13636364, 0.1355932 , 0.13483146, 0.1340782 ,\n",
              "        0.13333334, 0.13259669, 0.13186814, 0.13114755, 0.13043478,\n",
              "        0.12972972, 0.12903225, 0.12834224, 0.12765957, 0.12698412,\n",
              "        0.12631579, 0.12565446, 0.125     , 0.12435233, 0.12371133,\n",
              "        0.12307693, 0.12244899, 0.12182741, 0.12121212, 0.12060301],\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.42285714, 0.42045456, 0.41807905, 0.41573033, 0.4134078 ,\n",
              "        0.41111112, 0.4088398 , 0.4065934 , 0.40437162, 0.4021739 ,\n",
              "        0.39999998, 0.39784947, 0.39572194, 0.393617  , 0.3915344 ,\n",
              "        0.38947368, 0.38743457, 0.3854167 , 0.3834197 , 0.3814433 ,\n",
              "        0.3794872 , 0.37755105, 0.37563452, 0.37373737, 0.37185928],\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.70857143, 0.7045455 , 0.7005649 , 0.6966292 , 0.6927374 ,\n",
              "        0.6888889 , 0.6850829 , 0.6813187 , 0.6775957 , 0.67391306,\n",
              "        0.6702702 , 0.6666667 , 0.6631016 , 0.65957445, 0.65608466,\n",
              "        0.6526315 , 0.6492147 , 0.6458334 , 0.64248705, 0.63917524,\n",
              "        0.63589746, 0.6326531 , 0.6294416 , 0.6262626 , 0.6231156 ],\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",
              "       [0.9942857 , 0.9886364 , 0.98305076, 0.9775281 , 0.972067  ,\n",
              "        0.9666667 , 0.961326  , 0.95604396, 0.95081973, 0.9456522 ,\n",
              "        0.9405405 , 0.9354839 , 0.9304813 , 0.92553186, 0.92063487,\n",
              "        0.9157894 , 0.91099477, 0.90625   , 0.9015544 , 0.8969072 ,\n",
              "        0.8923077 , 0.88775516, 0.8832487 , 0.8787879 , 0.8743718 ],\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)"
            ]
          },
          "metadata": {},
          "execution_count": 110
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Global normalization using two axes at once"
      ],
      "metadata": {
        "id": "D4o8QRgXhoJP"
      }
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "eqZ8GBcFfttm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "narr = norm(arr)\n",
        "narr.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GKZKJTyKiBCy",
        "outputId": "2af26944-96f2-4cb5-e5df-4f6868e47ccb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 25)"
            ]
          },
          "metadata": {},
          "execution_count": 112
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(narr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Weeis4uNiGmI",
        "outputId": "f450ba48-7042-441e-8a7e-52d2cbb2cdf8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 113
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "narr"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NIc0LfUeiON_",
        "outputId": "3cecb9f0-dfbe-4731-915f-db354e0bc1f3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "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],\n",
              "       [1.2562814e-03, 1.3065326e-03, 1.3567839e-03, 1.4070352e-03,\n",
              "        1.4572864e-03, 1.5075377e-03, 1.5577889e-03, 1.6080402e-03,\n",
              "        1.6582914e-03, 1.7085427e-03, 1.7587940e-03, 1.8090452e-03,\n",
              "        1.8592965e-03, 1.9095477e-03, 1.9597989e-03, 2.0100502e-03,\n",
              "        2.0603016e-03, 2.1105527e-03, 2.1608039e-03, 2.2110553e-03,\n",
              "        2.2613066e-03, 2.3115578e-03, 2.3618089e-03, 2.4120603e-03,\n",
              "        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],\n",
              "       [3.7688441e-03, 3.8190954e-03, 3.8693468e-03, 3.9195977e-03,\n",
              "        3.9698491e-03, 4.0201005e-03, 4.0703518e-03, 4.1206032e-03,\n",
              "        4.1708541e-03, 4.2211055e-03, 4.2713569e-03, 4.3216078e-03,\n",
              "        4.3718591e-03, 4.4221105e-03, 4.4723619e-03, 4.5226132e-03,\n",
              "        4.5728642e-03, 4.6231155e-03, 4.6733669e-03, 4.7236178e-03,\n",
              "        4.7738692e-03, 4.8241206e-03, 4.8743719e-03, 4.9246233e-03,\n",
              "        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],\n",
              "       [6.2814071e-03, 6.3316580e-03, 6.3819094e-03, 6.4321607e-03,\n",
              "        6.4824121e-03, 6.5326635e-03, 6.5829144e-03, 6.6331658e-03,\n",
              "        6.6834171e-03, 6.7336680e-03, 6.7839194e-03, 6.8341708e-03,\n",
              "        6.8844222e-03, 6.9346735e-03, 6.9849244e-03, 7.0351758e-03,\n",
              "        7.0854272e-03, 7.1356781e-03, 7.1859295e-03, 7.2361808e-03,\n",
              "        7.2864322e-03, 7.3366836e-03, 7.3869345e-03, 7.4371858e-03,\n",
              "        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],\n",
              "       [8.7939696e-03, 8.8442210e-03, 8.8944724e-03, 8.9447238e-03,\n",
              "        8.9949751e-03, 9.0452265e-03, 9.0954769e-03, 9.1457283e-03,\n",
              "        9.1959797e-03, 9.2462311e-03, 9.2964824e-03, 9.3467338e-03,\n",
              "        9.3969852e-03, 9.4472356e-03, 9.4974870e-03, 9.5477384e-03,\n",
              "        9.5979897e-03, 9.6482411e-03, 9.6984925e-03, 9.7487438e-03,\n",
              "        9.7989952e-03, 9.8492466e-03, 9.8994970e-03, 9.9497484e-03,\n",
              "        9.9999998e-03]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 114
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Nested pmap"
      ],
      "metadata": {
        "id": "Cs_oFSYjkaIN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(8)).reshape(2,4)\n",
        "arr"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W-dMX9EckckK",
        "outputId": "0227f267-32c5-4786-91e7-c68955e12492"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0, 1, 2, 3],\n",
              "       [4, 5, 6, 7]], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 115
        }
      ]
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "YVEFRvlYkvcW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LVcHVBl9k7Tg",
        "outputId": "de33ffae-10e9-4c64-dd7e-ef0d4decfdee"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0.        , 0.03571429, 0.07142857, 0.10714287],\n",
              "       [0.14285715, 0.17857143, 0.21428573, 0.25      ]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 117
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(n(arr))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gAVypn8Lk-Ov",
        "outputId": "378a2b68-e69b-4441-bfb0-ae035fe303af"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 118
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The same example using the decorator style"
      ],
      "metadata": {
        "id": "vY_seo4knd9M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from functools import partial"
      ],
      "metadata": {
        "id": "RiikZf-lni6V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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'))"
      ],
      "metadata": {
        "id": "uBtUfQgonjer"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(n(arr))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XQeCcTBxnwdu",
        "outputId": "92353719-280a-4e74-fcd6-85f7a798efc7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 121
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data-parallel neural network training"
      ],
      "metadata": {
        "id": "DVXNIzeGNkaF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Preparing data"
      ],
      "metadata": {
        "id": "K0voYQTnSloK"
      }
    },
    {
      "cell_type": "code",
      "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']"
      ],
      "metadata": {
        "id": "M3B-FAZOOKOH",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1f815f28-b3f6-440d-d237-aa62c51e106c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/home/grigo/.local/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1mDownloading and preparing dataset 11.06 MiB (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /tmp/tfds/mnist/3.0.1...\u001b[0m\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Dl Completed...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 13.64 file/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1mDataset mnist downloaded and prepared to /tmp/tfds/mnist/3.0.1. Subsequent calls will reuse this data.\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "GQW9KLknOSc4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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",
        ")"
      ],
      "metadata": {
        "id": "XEL-NZOjOV0_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(train_data)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bcwCdU2-GU80",
        "outputId": "b303530c-c59a-4a87-ecb3-58ce29efda04"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "235"
            ]
          },
          "metadata": {},
          "execution_count": 129
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Preparing MLP"
      ],
      "metadata": {
        "id": "ilU3LzAeSolR"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "Q7B2PFjwOben"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "PARAM_SCALE = 0.01"
      ],
      "metadata": {
        "id": "aG1Ixe95OiIv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "id": "DIwXobuPOprS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "DjyD2_EmNlC5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Loss and update functions"
      ],
      "metadata": {
        "id": "TdwS4P06Sr-P"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "INIT_LR = 1.0\n",
        "DECAY_RATE = 0.95\n",
        "DECAY_STEPS = 5\n",
        "NUM_EPOCHS  = 20"
      ],
      "metadata": {
        "id": "bwg667wGO3Yy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from functools import partial"
      ],
      "metadata": {
        "id": "T7vWgrxUKqkH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def loss(params, images, targets):\n",
        "  \"\"\"Categorical cross entropy loss function.\"\"\"\n",
        "  logits = batched_predict(params, images)\n",
        "  log_preds = logits - logsumexp(logits) # 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"
      ],
      "metadata": {
        "id": "0V3nOttiO7da"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Section for debugging purposes"
      ],
      "metadata": {
        "id": "GKMrN2vOPLiW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_data_iter = iter(train_data)\n",
        "x, y = next(train_data_iter)"
      ],
      "metadata": {
        "id": "9l5wT9U0O-7J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x.shape, y.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fm2cdhekRNFC",
        "outputId": "9f02cf6b-02ec-4db1-988a-499fd0073b10"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((256, 28, 28, 1), (256,))"
            ]
          },
          "metadata": {},
          "execution_count": 138
        }
      ]
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MpJ5tnZ4PQMN",
        "outputId": "86579337-bd23-4e49-a5c2-eeeaffc47810"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 32, 784), (8, 32, 10))"
            ]
          },
          "metadata": {},
          "execution_count": 139
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "updated_params, loss_value = update(init_params, x, y, 0)"
      ],
      "metadata": {
        "id": "vSKemd9_QW1_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loss_value"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E3VpMr--Rf67",
        "outputId": "71f14b13-5e80-40f3-f015-06834caa845a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0.5771865 , 0.5766423 , 0.5766001 , 0.57689124, 0.57701343,\n",
              "       0.57676095, 0.57668227, 0.5764269 ], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 141
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training loop"
      ],
      "metadata": {
        "id": "4PFlq7EcS3q_"
      }
    },
    {
      "cell_type": "code",
      "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))"
      ],
      "metadata": {
        "id": "Zw4DIGl5InbS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "skYy98qTFUpP",
        "outputId": "3912fd05-dbf0-47e6-a46a-3650107a2192"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0 in 1.61 sec\n",
            "Training set loss 3.3886349201202393\n",
            "Training set accuracy 0.9229609370231628\n",
            "Test set accuracy 0.9217773675918579\n",
            "Epoch 1 in 1.01 sec\n",
            "Training set loss 3.0531487464904785\n",
            "Training set accuracy 0.9479831457138062\n",
            "Test set accuracy 0.9482421875\n",
            "Epoch 2 in 1.01 sec\n",
            "Training set loss 2.983154535293579\n",
            "Training set accuracy 0.9597517848014832\n",
            "Test set accuracy 0.9580078125\n",
            "Epoch 3 in 1.07 sec\n",
            "Training set loss 2.948392391204834\n",
            "Training set accuracy 0.9668937921524048\n",
            "Test set accuracy 0.96484375\n",
            "Epoch 4 in 1.07 sec\n",
            "Training set loss 2.928722858428955\n",
            "Training set accuracy 0.9714649319648743\n",
            "Test set accuracy 0.9706054925918579\n",
            "Epoch 5 in 0.95 sec\n",
            "Training set loss 2.913074254989624\n",
            "Training set accuracy 0.9739250540733337\n",
            "Test set accuracy 0.973339855670929\n",
            "Epoch 6 in 0.98 sec\n",
            "Training set loss 2.9012343883514404\n",
            "Training set accuracy 0.9771497845649719\n",
            "Test set accuracy 0.975292980670929\n",
            "Epoch 7 in 1.13 sec\n",
            "Training set loss 2.891705274581909\n",
            "Training set accuracy 0.9794880151748657\n",
            "Test set accuracy 0.9751953482627869\n",
            "Epoch 8 in 0.98 sec\n",
            "Training set loss 2.8835349082946777\n",
            "Training set accuracy 0.9813164472579956\n",
            "Test set accuracy 0.9766601920127869\n",
            "Epoch 9 in 0.97 sec\n",
            "Training set loss 2.8766467571258545\n",
            "Training set accuracy 0.9828124642372131\n",
            "Test set accuracy 0.97705078125\n",
            "Epoch 10 in 1.04 sec\n",
            "Training set loss 2.8705313205718994\n",
            "Training set accuracy 0.9838264584541321\n",
            "Test set accuracy 0.9784179925918579\n",
            "Epoch 11 in 1.00 sec\n",
            "Training set loss 2.8653604984283447\n",
            "Training set accuracy 0.9851396083831787\n",
            "Test set accuracy 0.97900390625\n",
            "Epoch 12 in 0.95 sec\n",
            "Training set loss 2.8611690998077393\n",
            "Training set accuracy 0.986120343208313\n",
            "Test set accuracy 0.97900390625\n",
            "Epoch 13 in 0.96 sec\n",
            "Training set loss 2.8567447662353516\n",
            "Training set accuracy 0.9867187142372131\n",
            "Test set accuracy 0.9791015982627869\n",
            "Epoch 14 in 1.01 sec\n",
            "Training set loss 2.853097438812256\n",
            "Training set accuracy 0.9877492785453796\n",
            "Test set accuracy 0.97998046875\n",
            "Epoch 15 in 0.92 sec\n",
            "Training set loss 2.8497848510742188\n",
            "Training set accuracy 0.9884640574455261\n",
            "Test set accuracy 0.9810547232627869\n",
            "Epoch 16 in 1.03 sec\n",
            "Training set loss 2.8475968837738037\n",
            "Training set accuracy 0.9889128804206848\n",
            "Test set accuracy 0.9810547232627869\n",
            "Epoch 17 in 1.00 sec\n",
            "Training set loss 2.8452184200286865\n",
            "Training set accuracy 0.9895278811454773\n",
            "Test set accuracy 0.980761706829071\n",
            "Epoch 18 in 0.95 sec\n",
            "Training set loss 2.8419547080993652\n",
            "Training set accuracy 0.9902426600456238\n",
            "Test set accuracy 0.981152355670929\n",
            "Epoch 19 in 1.00 sec\n",
            "Training set loss 2.83945631980896\n",
            "Training set accuracy 0.9909241795539856\n",
            "Test set accuracy 0.9814453125\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "zacBQRSbNSkB"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}