Appendix-D/JAX_in_Action_Appendix_D_pjit.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCirjpcFSPan"
      },
      "source": [
        "## ⚖️ Choose A or C:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5lNqNEGtLyO"
      },
      "source": [
        "## A: Emulating multi-device system on CPU\n",
        "\n",
        "Use this section to initialize a set of virtual devices on CPU if you have no access to a multi-device system.\n",
        "\n",
        "It can also help you prototype, debug and test your multi-device code locally before running it on the expensive system.\n",
        "\n",
        "Even in the case of using Google Colab it can help you prototype faster because a CPU runtime is faster to restart."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ypaH8OgftR_H"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "n8ip1VKgterO"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tlh75iWutfVT",
        "outputId": "d8efe2cf-4cc4-4f88-8225-058910645121"
      },
      "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
        }
      ],
      "source": [
        "jax.devices()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rmk_qf6mKCZx"
      },
      "source": [
        "## (do not use) B: Setting up Colab TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RMq70OC9wVHQ"
      },
      "source": [
        "#### [this section creates Colab TPU, which does not work with recent JAX versions anymore. Do not use TPU Runtime from Colab with this notebook]\n",
        "\n",
        "Use this section if you want to use Google Cloud TPU (and don't forget to change the Runtime type in \"Runtime\"-> \"Change runtime type\" -> \"TPU\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q9_o0MwOu94x"
      },
      "outputs": [],
      "source": [
        "# in order to use TPU you have to run this cell before importing JAX\n",
        "import jax.tools.colab_tpu\n",
        "jax.tools.colab_tpu.setup_tpu()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2aBy_QGmu-61",
        "outputId": "371f08da-64ff-4c18-a7d3-1db049cabae6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tpu\n"
          ]
        }
      ],
      "source": [
        "from jax.lib import xla_bridge\n",
        "print(xla_bridge.get_backend().platform)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7c2IXWjRvFbj"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_qLo4vNRvHc-",
        "outputId": "84ea03f5-1d4c-46c1-f73e-89e7b5f93da7"
      },
      "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": 4
        }
      ],
      "source": [
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## C: Using Cloud TPU and a Local runtime in Colab"
      ],
      "metadata": {
        "id": "pXcZUT7WJv4e"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make preparations according to the following links:\n",
        "\n",
        "* Creating a Cloud TPU https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm#tpu-vms\n",
        "\n",
        "* Preparing Jupyter and connect to a Local runtime https://research.google.com/colaboratory/local-runtimes.html\n"
      ],
      "metadata": {
        "id": "ps5o9t18KCb3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DECNlXiNKuc2",
        "outputId": "a2f7ff00-8d5c-42df-c739-b146bef18542"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in links: https://storage.googleapis.com/jax-releases/libtpu_releases.html\r\n",
            "Requirement already satisfied: jax[tpu] in ./.local/lib/python3.8/site-packages (0.4.13)\n",
            "Requirement already satisfied: ml-dtypes>=0.1.0 in ./.local/lib/python3.8/site-packages (from jax[tpu]) (0.2.0)\n",
            "Requirement already satisfied: numpy>=1.21 in ./.local/lib/python3.8/site-packages (from jax[tpu]) (1.24.3)\n",
            "Requirement already satisfied: opt-einsum in ./.local/lib/python3.8/site-packages (from jax[tpu]) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.7 in ./.local/lib/python3.8/site-packages (from jax[tpu]) (1.10.1)\n",
            "Requirement already satisfied: importlib-metadata>=4.6; python_version < \"3.10\" in ./.local/lib/python3.8/site-packages (from jax[tpu]) (6.8.0)\n",
            "Requirement already satisfied: jaxlib==0.4.13; extra == \"tpu\" in ./.local/lib/python3.8/site-packages (from jax[tpu]) (0.4.13)\n",
            "Requirement already satisfied: libtpu-nightly==0.1.dev20230622; extra == \"tpu\" in ./.local/lib/python3.8/site-packages (from jax[tpu]) (0.1.dev20230622)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/lib/python3/dist-packages (from importlib-metadata>=4.6; python_version < \"3.10\"->jax[tpu]) (1.0.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "9z7CUfFvJ7JN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.local_devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_RJ5XZCcKAYD",
        "outputId": "473d8adf-e348-4c4f-a25f-1bf75cc42c82"
      },
      "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",
      "metadata": {
        "id": "XRLNPE2Ese1D"
      },
      "source": [
        "## Using pjit()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "2hICp5VhC6Bt",
        "outputId": "825af820-1de1-4ee5-f3e4-a8e5111ea8b2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.4.23'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "jax.__version__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "e3HsZ2PIsl3t"
      },
      "outputs": [],
      "source": [
        "from jax.experimental.pjit import pjit"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "-Bocy386sr3n"
      },
      "outputs": [],
      "source": [
        "from jax import random"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1wKd0cB8iu5"
      },
      "source": [
        "### Old vmap+pmap example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "2p2R6uG5tVpV"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "SDMWf5eIswdB"
      },
      "outputs": [],
      "source": [
        "rng_key = random.PRNGKey(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "THq8XQoAs8CH",
        "outputId": "77d15cdc-5bcf-47ab-fb24-9db31858a9ca"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((10000000, 3), (10000000, 3))"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ],
      "source": [
        "vs = random.normal(rng_key, shape=(20_000_000,3))\n",
        "v1s = vs[:10_000_000,:]\n",
        "v2s = vs[10_000_000:,:]\n",
        "\n",
        "v1s.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jLM8V0BstMSD",
        "outputId": "4f9250ca-bbdc-4f3b-c5ad-c01493e098ac"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((8, 1250000, 3), (8, 1250000, 3))"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ],
      "source": [
        "v1sp = v1s.reshape((8, v1s.shape[0]//8, v1s.shape[1]))\n",
        "v2sp = v2s.reshape((8, v2s.shape[0]//8, v2s.shape[1]))\n",
        "\n",
        "v1sp.shape, v2sp.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "xf9UVhJcNHH4"
      },
      "outputs": [],
      "source": [
        "dot_parallel = jax.pmap(\n",
        "    jax.vmap(dot, in_axes=(0,0)),\n",
        "    in_axes=(0,0)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "8rdtr2F_NI_g"
      },
      "outputs": [],
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dF37XSrzNLE0",
        "outputId": "b861bba8-0d37-4e1d-a073-0a60a7a66e26"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000)"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "spEI1WDaNPXs",
        "outputId": "d4f243eb-49ed-49a6-9520-8342c7cf624e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gUeiCdt_uUmv"
      },
      "source": [
        "### Replacing with pjit and 1D mesh"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "poE9ZsTevKdY"
      },
      "outputs": [],
      "source": [
        "from jax.sharding import Mesh\n",
        "from jax.sharding import PartitionSpec"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "SI0alZjP56SZ"
      },
      "outputs": [],
      "source": [
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "26EiJI9V5vUA",
        "outputId": "2bd90481-1ead-4469-a97c-95e1779115f4"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3),\n",
              "       CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)],\n",
              "      dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "devices = np.array(jax.devices())\n",
        "devices"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_9oP_tvkRIrI"
      },
      "source": [
        "Non-vectorized function won't work here because output partitioning works only for rank>=1 tensors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "dR71WDupBiyb"
      },
      "outputs": [],
      "source": [
        "f = pjit(dot,\n",
        "         in_shardings =None,\n",
        "         out_shardings =PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 529
        },
        "id": "-fYCe7Q2BlSl",
        "outputId": "83f40d0e-f5c5-4d56-a04a-85d5429237b6"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "One of pjit outputs is incompatible with its sharding annotation NamedSharding(mesh=Mesh('devices': 8), spec=PartitionSpec('devices',)): Sharding NamedSharding(mesh=Mesh('devices': 8), spec=PartitionSpec('devices',)) is only valid for values of rank at least 1, but was applied to a value of rank 0. For scalars the PartitionSpec should be P()",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/sharding_impls.py\u001b[0m in \u001b[0;36mis_compatible_aval\u001b[0;34m(self, aval_shape)\u001b[0m\n\u001b[1;32m    282\u001b[0m                    if len(aval_shape) == 0 else '')\n\u001b[0;32m--> 283\u001b[0;31m       raise ValueError(\n\u001b[0m\u001b[1;32m    284\u001b[0m           \u001b[0;34mf\"Sharding {self} is only valid for values of rank at least \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: Sharding NamedSharding(mesh=Mesh('devices': 8), spec=PartitionSpec('devices',)) is only valid for values of rank at least 1, but was applied to a value of rank 0. For scalars the PartitionSpec should be P()",
            "\nDuring handling of the above exception, another exception occurred:\n",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-20-0153214dcc4d>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mMesh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevices\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'devices'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0mx_pjit\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv1s\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mv2s\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/pjit.py\u001b[0m in \u001b[0;36mpjit_check_aval_sharding\u001b[0;34m(shardings, flat_avals, names, what_aval, allow_uneven_sharding)\u001b[0m\n\u001b[1;32m   1021\u001b[0m         \u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_to_xla_hlo_sharding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1022\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mValueError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1023\u001b[0;31m       raise ValueError(\n\u001b[0m\u001b[1;32m   1024\u001b[0m           \u001b[0;34mf'One of {what_aval}{name_str} is incompatible with its sharding '\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1025\u001b[0m           f'annotation {s}: {e}')\n",
            "\u001b[0;31mValueError\u001b[0m: One of pjit outputs is incompatible with its sharding annotation NamedSharding(mesh=Mesh('devices': 8), spec=PartitionSpec('devices',)): Sharding NamedSharding(mesh=Mesh('devices': 8), spec=PartitionSpec('devices',)) is only valid for values of rank at least 1, but was applied to a value of rank 0. For scalars the PartitionSpec should be P()"
          ]
        }
      ],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yIfhsK4RRXCQ"
      },
      "source": [
        "Using a vectorized function that produces rank 1 tensor output. Input is replicated across all the devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "oy6ZX8j9_eVF"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=None,\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "fx46YXLx_4A_"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nppWGEcvAK8w",
        "outputId": "849c2ae3-9c06-4c68-cf9a-00d03d40e6d9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O0J-UmI0Balw",
        "outputId": "811814eb-d126-4463-8223-47f2712e7164"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ],
      "source": [
        "jax.numpy.all(x_pjit == x_pmap)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rAe6pb2wRhha"
      },
      "source": [
        "Input is sharded across all the devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "5_Y5ydczBxWl"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=PartitionSpec('devices'),\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "VZAAqX0eB-KV"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a8Ip8-pvCBKv",
        "outputId": "e09f2c4c-4f3b-44d9-ac3f-479fa4933f01"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SNxNE_yIRDYp",
        "outputId": "7854e396-0c50-4e7c-9ca7-79f22d6b4e95"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "jax.numpy.all(x_pjit == x_pmap)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LpH7bHKvaVUV"
      },
      "source": [
        "The same as previous"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "j_8SmhxM2rcQ"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=(PartitionSpec('devices'), None),\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "O7Cm1-hD2sjt"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EcMTZ_jZ2wzN",
        "outputId": "387225eb-0ed5-4248-ef93-92bf0f7e4cd5"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i_jJCF2JaX8n"
      },
      "source": [
        "Also the same"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "wXwWwPVQ5uJ6"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=PartitionSpec('devices', None),\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "_IA1QuIb5xgy"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "crGAiba_53Sw",
        "outputId": "863e829b-66f7-4566-c51e-fdf42bfe98da"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qbUs_y6GaaAj"
      },
      "source": [
        "Trying more values in PartitionSpec than there are parameters in the function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "1c6XEEa05KgF"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=(PartitionSpec('devices'), None, None),\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 366
        },
        "id": "quKUJlda5LiU",
        "outputId": "8b0494a9-ae05-423f-a30c-dca1acbb19fc"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "pjit in_shardings specification must be a tree prefix of the positional arguments tuple passed to the `pjit`-decorated function. In particular, pjit in_shardings must either be a None, a PartitionSpec, or a tuple of length equal to the number of positional arguments. But pjit in_shardings is the wrong length: got a tuple or list of length 3 for an args tuple of length 2.",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-36-0153214dcc4d>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mMesh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevices\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'devices'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m   \u001b[0mx_pjit\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv1s\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mv2s\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "    \u001b[0;31m[... skipping hidden 6 frame]\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/jax/_src/pjit.py\u001b[0m in \u001b[0;36mflatten_axis_resources\u001b[0;34m(what, tree, shardings, tupled_args)\u001b[0m\n\u001b[1;32m    874\u001b[0m                 f\"a singleton tuple.\")\n\u001b[1;32m    875\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 876\u001b[0;31m     \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    877\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    878\u001b[0m   \u001b[0maxis_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mshardings\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: pjit in_shardings specification must be a tree prefix of the positional arguments tuple passed to the `pjit`-decorated function. In particular, pjit in_shardings must either be a None, a PartitionSpec, or a tuple of length equal to the number of positional arguments. But pjit in_shardings is the wrong length: got a tuple or list of length 3 for an args tuple of length 2."
          ]
        }
      ],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ge3nElSNagye"
      },
      "source": [
        "Finally sharding both input parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "Ep0i7qNpTIVI"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=(PartitionSpec('devices'), PartitionSpec('devices')),\n",
        "         out_shardings=PartitionSpec('devices')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "id": "kbgIHp1ZTM9I"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('devices',)):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dpTYrHoBTNWI",
        "outputId": "79b5fb9e-5018-4e15-f83e-a44e490f0623"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UADUKdM2aM-N"
      },
      "source": [
        "### 2D mesh case"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "id": "ARSVrRjtu0si"
      },
      "outputs": [],
      "source": [
        "from jax.sharding import PartitionSpec as P   # could be useful to reduce typing\n",
        "from jax.sharding import Mesh\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "HgfhAh2HbSHX"
      },
      "outputs": [],
      "source": [
        "rng_key = random.PRNGKey(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yy0yLNP96WVF",
        "outputId": "6620675a-0cb9-410c-a14a-c0173d3a6d4b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((4000, 10000), (4000, 10000))"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ],
      "source": [
        "vs = random.normal(rng_key, shape=(8_000,10_000))\n",
        "v1s = vs[:4_000,:]\n",
        "v2s = vs[4_000:,:]\n",
        "\n",
        "v1s.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ePFmq1iObMoJ",
        "outputId": "47853f94-e7b4-4ada-b5da-2c3a23c2162f"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2),\n",
              "        CpuDevice(id=3)],\n",
              "       [CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6),\n",
              "        CpuDevice(id=7)]], dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ],
      "source": [
        "devices = np.array(jax.devices()).reshape(2, 4)\n",
        "devices"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "id": "UzfH0x9VbTLH"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "xC7XX6Ebbv7e"
      },
      "outputs": [],
      "source": [
        "f = pjit(jax.vmap(dot),\n",
        "         in_shardings=(P('x', 'y'), P('x', 'y')),\n",
        "         out_shardings=P('x')\n",
        "         )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "id": "bGS78ktTcgNh"
      },
      "outputs": [],
      "source": [
        "with Mesh(devices, ('x','y')):\n",
        "  x_pjit=f(v1s,v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sh_riKu2cj_A",
        "outputId": "c7d1b489-1f36-459d-b3c6-90370ad66832"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ],
      "source": [
        "x_pjit.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Je8_AQcdxYst"
      },
      "source": [
        "### Looking into HLO"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F5rmPtoWxgtJ"
      },
      "source": [
        "You can't see any collective ops on the jaxpr level:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_UfVcZ6Te67x",
        "outputId": "f013cb73-3180-4007-a206-652708412cf1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{ lambda ; a:f32[4000,10000] b:f32[4000,10000]. let\n",
            "    c:f32[4000] = pjit[\n",
            "      name=dot\n",
            "      in_shardings=(GSPMDSharding({devices=[2,4]<=[8]}), GSPMDSharding({devices=[2,4]<=[8]}))\n",
            "      jaxpr={ lambda ; d:f32[4000,10000] e:f32[4000,10000]. let\n",
            "          f:f32[4000] = dot_general[\n",
            "            dimension_numbers=(([1], [1]), ([0], [0]))\n",
            "            preferred_element_type=float32\n",
            "          ] d e\n",
            "        in (f,) }\n",
            "      out_shardings=(GSPMDSharding({devices=[2,4]<=[8] last_tile_dim_replicate}),)\n",
            "      resource_env=ResourceEnv(mesh=Mesh('x': 2, 'y': 4), ())\n",
            "    ] a b\n",
            "  in (c,) }\n"
          ]
        }
      ],
      "source": [
        "with Mesh(devices, ('x','y')):\n",
        "  print(jax.make_jaxpr(f)(v1s,v2s))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WU0c2nGbxlUi"
      },
      "source": [
        "But you can see them after the compilation:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "with Mesh(devices, ('x','y')):\n",
        "   print(f.lower(v1s, v2s).compile().as_text())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rgfwV61EAoUz",
        "outputId": "0e4e3eba-56d2-466b-be6a-b18ef48f96d1"
      },
      "execution_count": 65,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule pjit_dot, entry_computation_layout={(f32[2000,2500]{1,0}, f32[2000,2500]{1,0})->f32[2000]{0}}, num_partitions=8\n",
            "\n",
            "%add.clone (x.1: f32[], y.1: f32[]) -> f32[] {\n",
            "  %x.1 = f32[] parameter(0)\n",
            "  %y.1 = f32[] parameter(1)\n",
            "  ROOT %add.1 = f32[] add(f32[] %x.1, f32[] %y.1)\n",
            "}\n",
            "\n",
            "ENTRY %main.6_spmd (param: f32[2000,2500], param.1: f32[2000,2500]) -> f32[2000] {\n",
            "  %param = f32[2000,2500]{1,0} parameter(0), sharding={devices=[2,4]<=[8]}\n",
            "  %param.1 = f32[2000,2500]{1,0} parameter(1), sharding={devices=[2,4]<=[8]}\n",
            "  %dot = f32[2000]{0} dot(f32[2000,2500]{1,0} %param, f32[2000,2500]{1,0} %param.1), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}, metadata={op_name=\"pjit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=float32]\" source_file=\"<ipython-input-44-ccdf370c8172>\" source_line=2}\n",
            "  ROOT %all-reduce = f32[2000]{0} all-reduce(f32[2000]{0} %dot), channel_id=1, replica_groups={{0,1,2,3},{4,5,6,7}}, use_global_device_ids=true, to_apply=%add.clone, metadata={op_name=\"pjit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=float32]\" source_file=\"<ipython-input-44-ccdf370c8172>\" source_line=2}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 245
        },
        "id": "yCgBvwYPr6yA",
        "outputId": "9580f0c7-8682-452b-c7c3-7ed0aed3560c"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "AttributeError",
          "evalue": "'Compiled' object has no attribute 'compiler_ir'",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-66-beaa777c89ec>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Thanks to https://github.com/google/jax/discussions/11275\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mMesh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevices\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'x'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'y'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0mmodules\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlower\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv1s\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv2s\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompiler_ir\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mhlo\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mmodules\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhlo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_string\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mAttributeError\u001b[0m: 'Compiled' object has no attribute 'compiler_ir'"
          ]
        }
      ],
      "source": [
        "# (doesn't work anymore)\n",
        "# Thanks to https://github.com/google/jax/discussions/11275\n",
        "with Mesh(devices, ('x','y')):\n",
        "    modules = f.lower(v1s, v2s).compile().compiler_ir()\n",
        "    for hlo in modules:\n",
        "        print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uwkURlR91YHR"
      },
      "source": [
        "## MLP example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NZySdFPR1fWr"
      },
      "source": [
        "### Preparing data"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install these modules if you created a new empty cloud machine"
      ],
      "metadata": {
        "id": "wFW5RFD2NqOI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lvgUueADKium",
        "outputId": "9df941dd-cf07-404d-bcad-676f6bb668d0"
      },
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorflow in /usr/local/lib/python3.10/dist-packages (2.15.0)\n",
            "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.4.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.6.3)\n",
            "Requirement already satisfied: flatbuffers>=23.5.26 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (23.5.26)\n",
            "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (0.5.4)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (0.2.0)\n",
            "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (3.9.0)\n",
            "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (16.0.6)\n",
            "Requirement already satisfied: ml-dtypes~=0.2.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (0.2.0)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.23.5)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (3.3.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorflow) (23.2)\n",
            "Requirement already satisfied: 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 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (3.20.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from tensorflow) (67.7.2)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.16.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (2.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (4.5.0)\n",
            "Requirement already satisfied: wrapt<1.15,>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.14.1)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (0.35.0)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (1.60.0)\n",
            "Requirement already satisfied: tensorboard<2.16,>=2.15 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (2.15.1)\n",
            "Requirement already satisfied: tensorflow-estimator<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (2.15.0)\n",
            "Requirement already satisfied: keras<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow) (2.15.0)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from astunparse>=1.6.0->tensorflow) (0.42.0)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (2.17.3)\n",
            "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (1.2.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (3.5.1)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (2.31.0)\n",
            "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (0.7.2)\n",
            "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow) (3.0.1)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (5.3.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (0.3.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (4.9)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow) (1.3.1)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow) (2023.11.17)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=1.0.1->tensorboard<2.16,>=2.15->tensorflow) (2.1.3)\n",
            "Requirement already satisfied: pyasn1<0.6.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow) (0.5.1)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow) (3.2.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow_datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R1PpLK59K12_",
        "outputId": "3fbc1dda-489b-4672-9bfb-2e5759448f16"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorflow_datasets in /usr/local/lib/python3.10/dist-packages (4.9.4)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (1.4.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (8.1.7)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (0.1.8)\n",
            "Requirement already satisfied: etils[enp,epath,etree]>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (1.6.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (1.23.5)\n",
            "Requirement already satisfied: promise in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (2.3)\n",
            "Requirement already satisfied: protobuf>=3.20 in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (3.20.3)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (5.9.5)\n",
            "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (2.31.0)\n",
            "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (1.14.0)\n",
            "Requirement already satisfied: termcolor in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (2.4.0)\n",
            "Requirement already satisfied: toml in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (0.10.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (4.66.1)\n",
            "Requirement already satisfied: wrapt in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (1.14.1)\n",
            "Requirement already satisfied: array-record>=0.5.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow_datasets) (0.5.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,etree]>=0.9.0->tensorflow_datasets) (2023.6.0)\n",
            "Requirement already satisfied: importlib_resources in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,etree]>=0.9.0->tensorflow_datasets) (6.1.1)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,etree]>=0.9.0->tensorflow_datasets) (4.5.0)\n",
            "Requirement already satisfied: zipp in /usr/local/lib/python3.10/dist-packages (from etils[enp,epath,etree]>=0.9.0->tensorflow_datasets) (3.17.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow_datasets) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow_datasets) (3.6)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2023.11.17)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from promise->tensorflow_datasets) (1.16.0)\n",
            "Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow-metadata->tensorflow_datasets) (1.62.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {
        "id": "BrWQqOqTuz8C",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 84,
          "referenced_widgets": [
            "c9f8cc60935c4f9e8aa7cc000ea05887",
            "d761428cd59f468484cc3527a89e3976",
            "22df43ea2a444224827221c29b5ea07b",
            "d49d8508c64841a58933ab41420aac8d",
            "631605f468b7471b9e58fa44ef0bf690",
            "738eafdfbd284043a6b60d9140d85de8",
            "f4980724e69142e389dd040be49845c4",
            "45bf9ab0c6264823ba040f3ab4201f0b",
            "b65993253bce49e9a1a308dc207f917a",
            "fa0cbc1c4e6d4b48bd503333e7e200b7",
            "1a42f88802f64171904bcc3126893989"
          ]
        },
        "outputId": "0d27c471-84f2-4b89-efce-467382378b03"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading 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...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Dl Completed...:   0%|          | 0/5 [00:00<?, ? file/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c9f8cc60935c4f9e8aa7cc000ea05887"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset mnist downloaded and prepared to /tmp/tfds/mnist/3.0.1. Subsequent calls will reuse this data.\n"
          ]
        }
      ],
      "source": [
        "import jax\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "data_dir = '/tmp/tfds'\n",
        "\n",
        "data, info = tfds.load(name=\"mnist\",\n",
        "                       data_dir=data_dir,\n",
        "                       as_supervised=True,\n",
        "                       with_info=True)\n",
        "\n",
        "data_train = data['train']\n",
        "data_test  = data['test']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {
        "id": "qEONTKl81iKa"
      },
      "outputs": [],
      "source": [
        "HEIGHT = 28\n",
        "WIDTH  = 28\n",
        "CHANNELS = 1\n",
        "NUM_PIXELS = HEIGHT * WIDTH * CHANNELS\n",
        "NUM_LABELS = info.features['label'].num_classes\n",
        "NUM_DEVICES = jax.device_count()\n",
        "BATCH_SIZE  = 32"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "id": "OwXDzEbT1j61"
      },
      "outputs": [],
      "source": [
        "def preprocess(img, label):\n",
        "  \"\"\"Resize and preprocess images.\"\"\"\n",
        "  return (tf.cast(img, tf.float32)/255.0), label\n",
        "\n",
        "train_data = tfds.as_numpy(\n",
        "    data_train.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")\n",
        "test_data  = tfds.as_numpy(\n",
        "    data_test.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "seqmmqqx1let",
        "outputId": "32f9fbb6-4749-4959-ac36-13cd238ff8c1"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "235"
            ]
          },
          "metadata": {},
          "execution_count": 72
        }
      ],
      "source": [
        "len(train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TkQT5al_1ove"
      },
      "source": [
        "### Preparing MLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "id": "ajNo9yJY1mwQ"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jax import grad, jit, vmap, value_and_grad\n",
        "from jax import random\n",
        "from jax.nn import swish, logsumexp, one_hot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "id": "oRrZN0Xh1qsl"
      },
      "outputs": [],
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "PARAM_SCALE = 0.01"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "id": "SjxCJJLD1sEE"
      },
      "outputs": [],
      "source": [
        "def init_network_params(sizes, key=random.PRNGKey(0), scale=1e-2):\n",
        "  \"\"\"Initialize all layers for a fully-connected neural network with given sizes\"\"\"\n",
        "\n",
        "  def random_layer_params(m, n, key, scale=1e-2):\n",
        "    \"\"\"A helper function to randomly initialize weights and biases of a dense layer\"\"\"\n",
        "    w_key, b_key = random.split(key)\n",
        "    return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n",
        "\n",
        "  keys = random.split(key, len(sizes))\n",
        "  return [random_layer_params(m, n, k, scale) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n",
        "\n",
        "init_params = init_network_params(LAYER_SIZES, random.PRNGKey(0), scale=PARAM_SCALE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "id": "Kl6UEZ8d1t-s"
      },
      "outputs": [],
      "source": [
        "def predict(params, image):\n",
        "  \"\"\"Function for per-example predictions.\"\"\"\n",
        "  activations = image\n",
        "  for w, b in params[:-1]:\n",
        "    outputs = jnp.dot(w, activations) + b\n",
        "    activations = swish(outputs)\n",
        "\n",
        "  final_w, final_b = params[-1]\n",
        "  logits = jnp.dot(final_w, activations) + final_b\n",
        "  return logits\n",
        "\n",
        "batched_predict = vmap(predict, in_axes=(None, 0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Eo_z-kZ1yK6"
      },
      "source": [
        "### Loss and update functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "id": "MNYNpeuPTyYj"
      },
      "outputs": [],
      "source": [
        "from jax.experimental.pjit import pjit\n",
        "from jax.sharding import PartitionSpec as P\n",
        "from jax.sharding import Mesh\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "rzxpgkMj1vwl"
      },
      "outputs": [],
      "source": [
        "INIT_LR = 1.0\n",
        "DECAY_RATE = 0.95\n",
        "DECAY_STEPS = 5\n",
        "NUM_EPOCHS  = 20"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "id": "tPkJ5dYp11Q4"
      },
      "outputs": [],
      "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",
        "def update(params, x, y, epoch_number):\n",
        "  loss_value, grads = value_and_grad(loss)(params, x, y)\n",
        "  lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS)\n",
        "  return [(w - lr * dw, b - lr * db)\n",
        "          for (w, b), (dw, db) in zip(params, grads)], loss_value"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "id": "e9LhptZvT8TJ"
      },
      "outputs": [],
      "source": [
        "f_update = pjit(update,\n",
        "         in_shardings=(None, P('x'), P('x'), None),\n",
        "         out_shardings=None\n",
        "         )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XFyXVUe62Epe"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B65sfV72T6bT",
        "outputId": "fe05c997-6aeb-40e3-a44d-bb9d2e143740"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3),\n",
              "       CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)],\n",
              "      dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 81
        }
      ],
      "source": [
        "devices = np.array(jax.devices())\n",
        "devices"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {
        "id": "MjDtmztL1925"
      },
      "outputs": [],
      "source": [
        "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",
        "f_batch_accuracy = pjit(batch_accuracy,\n",
        "         in_shardings=(None, P('x'), P('x')),\n",
        "         out_shardings=None\n",
        "         )\n",
        "\n",
        "def accuracy(params, data):\n",
        "  accs = []\n",
        "  for images, targets in data:\n",
        "    accs.append(f_batch_accuracy(params, images, targets))\n",
        "  return jnp.mean(jnp.array(accs))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CGd4pczh2I1h",
        "outputId": "1c21fc13-af69-4b6e-b81b-20cfd6934f25"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0 in 10.87 sec\n",
            "Training set loss 0.691484808921814\n",
            "Training set accuracy 0.8647883534431458\n",
            "Test set accuracy 0.875\n",
            "Epoch 1 in 7.85 sec\n",
            "Training set loss 0.6260023713111877\n",
            "Training set accuracy 0.8881925940513611\n",
            "Test set accuracy 0.8935546875\n",
            "Epoch 2 in 7.83 sec\n",
            "Training set loss 0.6165075302124023\n",
            "Training set accuracy 0.8978723287582397\n",
            "Test set accuracy 0.9009765982627869\n",
            "Epoch 3 in 7.42 sec\n",
            "Training set loss 0.6100791692733765\n",
            "Training set accuracy 0.9065104126930237\n",
            "Test set accuracy 0.9090820550918579\n",
            "Epoch 4 in 7.19 sec\n",
            "Training set loss 0.6044540405273438\n",
            "Training set accuracy 0.9142231941223145\n",
            "Test set accuracy 0.917285144329071\n",
            "Epoch 5 in 7.66 sec\n",
            "Training set loss 0.5995851755142212\n",
            "Training set accuracy 0.9210605025291443\n",
            "Test set accuracy 0.924609363079071\n",
            "Epoch 6 in 6.53 sec\n",
            "Training set loss 0.5956109166145325\n",
            "Training set accuracy 0.9260915517807007\n",
            "Test set accuracy 0.9291015863418579\n",
            "Epoch 7 in 6.44 sec\n",
            "Training set loss 0.5924531817436218\n",
            "Training set accuracy 0.9303967356681824\n",
            "Test set accuracy 0.9320312738418579\n",
            "Epoch 8 in 6.81 sec\n",
            "Training set loss 0.5899131894111633\n",
            "Training set accuracy 0.9341201186180115\n",
            "Test set accuracy 0.9364258050918579\n",
            "Epoch 9 in 6.85 sec\n",
            "Training set loss 0.5878121852874756\n",
            "Training set accuracy 0.936995804309845\n",
            "Test set accuracy 0.938769519329071\n",
            "Epoch 10 in 6.60 sec\n",
            "Training set loss 0.5860452651977539\n",
            "Training set accuracy 0.9397329688072205\n",
            "Test set accuracy 0.942187488079071\n",
            "Epoch 11 in 6.48 sec\n",
            "Training set loss 0.5845457911491394\n",
            "Training set accuracy 0.9419769644737244\n",
            "Test set accuracy 0.944042980670929\n",
            "Epoch 12 in 6.84 sec\n",
            "Training set loss 0.5832619071006775\n",
            "Training set accuracy 0.9444869160652161\n",
            "Test set accuracy 0.9461914300918579\n",
            "Epoch 13 in 6.61 sec\n",
            "Training set loss 0.582151472568512\n",
            "Training set accuracy 0.9467143416404724\n",
            "Test set accuracy 0.947558581829071\n",
            "Epoch 14 in 6.52 sec\n",
            "Training set loss 0.5811813473701477\n",
            "Training set accuracy 0.9483765363693237\n",
            "Test set accuracy 0.948437511920929\n",
            "Epoch 15 in 6.94 sec\n",
            "Training set loss 0.5803254842758179\n",
            "Training set accuracy 0.9500720500946045\n",
            "Test set accuracy 0.94970703125\n",
            "Epoch 16 in 8.03 sec\n",
            "Training set loss 0.5795639157295227\n",
            "Training set accuracy 0.9512355923652649\n",
            "Test set accuracy 0.950976550579071\n",
            "Epoch 17 in 6.50 sec\n",
            "Training set loss 0.5788805484771729\n",
            "Training set accuracy 0.9528313279151917\n",
            "Test set accuracy 0.9527344107627869\n",
            "Epoch 18 in 6.93 sec\n",
            "Training set loss 0.5782625079154968\n",
            "Training set accuracy 0.9543938636779785\n",
            "Test set accuracy 0.9541015625\n",
            "Epoch 19 in 6.64 sec\n",
            "Training set loss 0.5776995420455933\n",
            "Training set accuracy 0.9556239247322083\n",
            "Test set accuracy 0.9554687738418579\n"
          ]
        }
      ],
      "source": [
        "import time\n",
        "\n",
        "params = init_params\n",
        "with Mesh(devices, ('x',)):\n",
        "  for epoch in range(NUM_EPOCHS):\n",
        "    start_time = time.time()\n",
        "    losses = []\n",
        "    for x, y in train_data:\n",
        "      x = jnp.reshape(x, (len(x), NUM_PIXELS))\n",
        "      y = one_hot(y, NUM_LABELS)\n",
        "      params, loss_value = f_update(params, x, y, epoch)\n",
        "      losses.append(jnp.sum(loss_value))\n",
        "    epoch_time = time.time() - start_time\n",
        "\n",
        "    train_acc = accuracy(params, train_data)\n",
        "    test_acc = accuracy(params, test_data)\n",
        "    print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
        "    print(\"Training set loss {}\".format(jnp.mean(jnp.array(losses))))\n",
        "    print(\"Training set accuracy {}\".format(train_acc))\n",
        "    print(\"Test set accuracy {}\".format(test_acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Some resources that might help you achieve weight sharding as well:\n",
        "https://github.com/google/jax/discussions/8649"
      ],
      "metadata": {
        "id": "CpzcEaL836k6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**!!! Do not forget to shutdown your Cloud TPU, or you'll spend much money on it!!!**"
      ],
      "metadata": {
        "id": "gJeoJJNGMrxA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0igAXyAE2dtf"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "c9f8cc60935c4f9e8aa7cc000ea05887": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_d761428cd59f468484cc3527a89e3976",
              "IPY_MODEL_22df43ea2a444224827221c29b5ea07b",
              "IPY_MODEL_d49d8508c64841a58933ab41420aac8d"
            ],
            "layout": "IPY_MODEL_631605f468b7471b9e58fa44ef0bf690"
          }
        },
        "d761428cd59f468484cc3527a89e3976": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_738eafdfbd284043a6b60d9140d85de8",
            "placeholder": "​",
            "style": "IPY_MODEL_f4980724e69142e389dd040be49845c4",
            "value": "Dl Completed...: 100%"
          }
        },
        "22df43ea2a444224827221c29b5ea07b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_45bf9ab0c6264823ba040f3ab4201f0b",
            "max": 5,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b65993253bce49e9a1a308dc207f917a",
            "value": 5
          }
        },
        "d49d8508c64841a58933ab41420aac8d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_fa0cbc1c4e6d4b48bd503333e7e200b7",
            "placeholder": "​",
            "style": "IPY_MODEL_1a42f88802f64171904bcc3126893989",
            "value": " 5/5 [00:00&lt;00:00,  9.05 file/s]"
          }
        },
        "631605f468b7471b9e58fa44ef0bf690": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "738eafdfbd284043a6b60d9140d85de8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f4980724e69142e389dd040be49845c4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "45bf9ab0c6264823ba040f3ab4201f0b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b65993253bce49e9a1a308dc207f917a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "fa0cbc1c4e6d4b48bd503333e7e200b7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1a42f88802f64171904bcc3126893989": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}