Chapter-8/JAX_in_Action_Chapter_8_Tensor_Sharding.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Using Cloud TPU and a Local runtime in Colab"
      ],
      "metadata": {
        "id": "pXcZUT7WJv4e"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make preparations according to Appendix C, or 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": {
        "id": "DECNlXiNKuc2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e8cf6e70-1357-496d-aa30-95506609e5b4"
      },
      "execution_count": 1,
      "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)\r\n",
            "Requirement already satisfied: ml-dtypes>=0.1.0 in ./.local/lib/python3.8/site-packages (from jax[tpu]) (0.2.0)\r\n",
            "Requirement already satisfied: numpy>=1.21 in ./.local/lib/python3.8/site-packages (from jax[tpu]) (1.24.4)\r\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]) (7.0.1)\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": "markdown",
      "source": [
        "Needed for pretty sharding visualization"
      ],
      "metadata": {
        "id": "-GXV0DH2P7kf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install rich"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ryHjnoA7P5pQ",
        "outputId": "6e9b1fc6-cefb-4d71-8447-b81ac60340f5"
      },
      "execution_count": 2,
      "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 rich\n",
            "  Downloading rich-13.7.0-py3-none-any.whl (240 kB)\n",
            "\u001b[K     |████████████████████████████████| 240 kB 1.3 MB/s \n",
            "\u001b[?25hCollecting markdown-it-py>=2.2.0\n",
            "  Downloading markdown_it_py-3.0.0-py3-none-any.whl (87 kB)\n",
            "\u001b[K     |████████████████████████████████| 87 kB 6.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pygments<3.0.0,>=2.13.0 in ./.local/lib/python3.8/site-packages (from rich) (2.17.2)\n",
            "Requirement already satisfied: typing-extensions<5.0,>=4.0.0; python_version < \"3.9\" in ./.local/lib/python3.8/site-packages (from rich) (4.9.0)\n",
            "Collecting mdurl~=0.1\n",
            "  Downloading mdurl-0.1.2-py3-none-any.whl (10.0 kB)\n",
            "Installing collected packages: mdurl, markdown-it-py, rich\n",
            "Successfully installed markdown-it-py-3.0.0 mdurl-0.1.2 rich-13.7.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "9z7CUfFvJ7JN"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.local_devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_RJ5XZCcKAYD",
        "outputId": "06470233-f45e-4085-8b53-7dafdfd10bfd"
      },
      "execution_count": 4,
      "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
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2hICp5VhC6Bt",
        "outputId": "ecb48f78-7eb3-4c36-b36e-9498defbee2a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.4.13'"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "jax.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XRLNPE2Ese1D"
      },
      "source": [
        "## Using Tensor Sharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "e3HsZ2PIsl3t"
      },
      "outputs": [],
      "source": [
        "from jax.experimental import mesh_utils\n",
        "from jax.sharding import PositionalSharding, NamedSharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "-Bocy386sr3n"
      },
      "outputs": [],
      "source": [
        "from jax import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "SI0alZjP56SZ"
      },
      "outputs": [],
      "source": [
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1wKd0cB8iu5"
      },
      "source": [
        "### Dot example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "2p2R6uG5tVpV"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "SDMWf5eIswdB"
      },
      "outputs": [],
      "source": [
        "rng_key = random.PRNGKey(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "THq8XQoAs8CH",
        "outputId": "4a2d4210-5c82-4aa0-a4f2-bee282d41628"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((4000, 10000), (4000, 10000))"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "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",
      "source": [
        "jax.debug.visualize_array_sharding(v1s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "ap2c7J2lPz-H",
        "outputId": "7064d552-647e-41e4-b39f-87b946dab584"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Positional sharding"
      ],
      "metadata": {
        "id": "we0gE4dV9bHv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((8,1)))"
      ],
      "metadata": {
        "id": "D6OxkmCbQeJw"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sharding"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C86DIH_uQ7DT",
        "outputId": "1b2a7d3e-08eb-41f9-90d0-c220296161d6"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 0}]\n",
              "                    [{TPU 1}]\n",
              "                    [{TPU 2}]\n",
              "                    [{TPU 3}]\n",
              "                    [{TPU 6}]\n",
              "                    [{TPU 7}]\n",
              "                    [{TPU 4}]\n",
              "                    [{TPU 5}]])"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = jax.device_put(v1s, sharding)"
      ],
      "metadata": {
        "id": "6YuoF1hhRgRs"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "type(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3ASNkCAddFkB",
        "outputId": "3a1ac310-78ac-4691-9b07-8d99271e7d65"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "jaxlib.xla_extension.ArrayImpl"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "DRm5TQDrYwqo",
        "outputId": "e6a3b99a-d696-47a4-e571-438028703898"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m1\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m2\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m3\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m6\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m7\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m4\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m5\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v2sp = jax.device_put(v2s, sharding)"
      ],
      "metadata": {
        "id": "9N2J_H_KRRc6"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "0_5ZNoMZQhrZ",
        "outputId": "b4cbcf66-f46b-43ae-b190-36e4edbb5f36"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m1\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m2\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m3\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m6\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m7\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m4\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m5\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rAe6pb2wRhha"
      },
      "source": [
        "Input is sharded across all the devices."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ],
      "metadata": {
        "id": "qjjU9AGsu9cX"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "d.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vnl0vIl3vAOc",
        "outputId": "9dadfa81-0d10-4561-c44b-4686a4317931"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "lM6bu5xqvFlu",
        "outputId": "2ef18cfd-c9a7-47e1-9f93-f9e64a2ae16c"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐\n",
              "│ TPU \u001b[1;36m0\u001b[0m │ TPU \u001b[1;36m1\u001b[0m │ TPU \u001b[1;36m2\u001b[0m │ TPU \u001b[1;36m3\u001b[0m │ TPU \u001b[1;36m6\u001b[0m │ TPU \u001b[1;36m7\u001b[0m │ TPU \u001b[1;36m4\u001b[0m │ TPU \u001b[1;36m5\u001b[0m │\n",
              "└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span> │\n",
              "└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jax.vmap(dot)(v1sp, v2sp).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r6LQ3H0xvV-S",
        "outputId": "e57a9bf2-bba4-439e-dc56-4602d8a88942"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.54 ms ± 26.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jax.vmap(dot)(v1s, v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rdbMpGTkvcrw",
        "outputId": "eaf3f7df-2a59-4aba-ccdb-b2462fa0c2ad"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.13 ms ± 51.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(jax.vmap(dot))(v1s, v2s)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ovKaom8nvoN-",
        "outputId": "8d6e80b7-3709-439d-cfb1-33f54407df95"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[4000,10000]\u001b[39m b\u001b[35m:f32[4000,10000]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[4000]\u001b[39m = dot_general[dimension_numbers=(([1], [1]), ([0], [0]))] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1s, v2s).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9MLQkQxBv0LP",
        "outputId": "66b16e90-5dac-46ad-e55d-c1f5c1e5298c"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[4000,10000]{1,0:T(8,128)}, f32[4000,10000]{1,0:T(8,128)})->f32[4000]{0:T(1024)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(256)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(256)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(256)} add(f32[]{:T(256)} %scalar_lhs, f32[]{:T(256)} %scalar_rhs)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.2: f32[4000,10000], param_1.2: f32[4000,10000]) -> f32[4000] {\n",
            "  %param_0.2 = f32[4000,10000]{1,0:T(8,128)} parameter(0)\n",
            "  %param_1.2 = f32[4000,10000]{1,0:T(8,128)} parameter(1)\n",
            "  %multiply.1 = f32[4000,10000]{1,0:T(8,128)} multiply(f32[4000,10000]{1,0:T(8,128)} %param_0.2, f32[4000,10000]{1,0:T(8,128)} %param_1.2)\n",
            "  %constant.1 = f32[]{:T(256)} constant(0)\n",
            "  ROOT %reduce.1 = f32[4000]{0:T(1024)} reduce(f32[4000,10000]{1,0:T(8,128)} %multiply.1, f32[]{:T(256)} %constant.1), dimensions={1}, to_apply=%scalar_add_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}\n",
            "}\n",
            "\n",
            "ENTRY %main.4 (Arg_0.1: f32[4000,10000], Arg_1.2: f32[4000,10000]) -> f32[4000] {\n",
            "  %Arg_1.2 = f32[4000,10000]{1,0:T(8,128)} parameter(1), sharding={replicated}\n",
            "  %Arg_0.1 = f32[4000,10000]{1,0:T(8,128)} parameter(0), sharding={replicated}\n",
            "  ROOT %fusion = f32[4000]{0:T(1024)} fusion(f32[4000,10000]{1,0:T(8,128)} %Arg_0.1, f32[4000,10000]{1,0:T(8,128)} %Arg_1.2), kind=kLoop, calls=%fused_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"16\",\"40\"],\"input_window_bounds\":[],\"estimated_cycles\":\"1125744\",\"iteration_bounds\":[\"32\",\"2\"]},\"scoped_memory_configs\":[]}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.make_jaxpr(jax.vmap(dot))(v1sp, v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HfVoi5Z72lUG",
        "outputId": "61032fcf-266a-4e1c-8e05-40ea521549e1"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{ \u001b[34m\u001b[22m\u001b[1mlambda \u001b[39m\u001b[22m\u001b[22m; a\u001b[35m:f32[4000,10000]\u001b[39m b\u001b[35m:f32[4000,10000]\u001b[39m. \u001b[34m\u001b[22m\u001b[1mlet\n",
              "    \u001b[39m\u001b[22m\u001b[22mc\u001b[35m:f32[4000]\u001b[39m = dot_general[dimension_numbers=(([1], [1]), ([0], [0]))] a b\n",
              "  \u001b[34m\u001b[22m\u001b[1min \u001b[39m\u001b[22m\u001b[22m(c,) }"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1sp, v2sp).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uRElWOlm2o57",
        "outputId": "564ada03-cd38-40e0-816c-0fc278bb84f1"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[500,10000]{1,0:T(8,128)}, f32[500,10000]{1,0:T(8,128)})->f32[500]{0:T(512)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(256)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(256)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(256)} add(f32[]{:T(256)} %scalar_lhs, f32[]{:T(256)} %scalar_rhs)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.2: f32[500,10000], param_1.2: f32[500,10000]) -> f32[500] {\n",
            "  %param_0.2 = f32[500,10000]{1,0:T(8,128)} parameter(0)\n",
            "  %param_1.2 = f32[500,10000]{1,0:T(8,128)} parameter(1)\n",
            "  %multiply.2 = f32[500,10000]{1,0:T(8,128)} multiply(f32[500,10000]{1,0:T(8,128)} %param_0.2, f32[500,10000]{1,0:T(8,128)} %param_1.2)\n",
            "  %constant.2 = f32[]{:T(256)} constant(0)\n",
            "  ROOT %reduce.2 = f32[500]{0:T(512)} reduce(f32[500,10000]{1,0:T(8,128)} %multiply.2, f32[]{:T(256)} %constant.2), dimensions={1}, to_apply=%scalar_add_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[500,10000], param.1: f32[500,10000]) -> f32[500] {\n",
            "  %param.1 = f32[500,10000]{1,0:T(8,128)} parameter(1), sharding={devices=[8,1]0,1,2,3,4,5,6,7}\n",
            "  %param = f32[500,10000]{1,0:T(8,128)} parameter(0), sharding={devices=[8,1]0,1,2,3,4,5,6,7}\n",
            "  ROOT %fusion = f32[500]{0:T(512)} fusion(f32[500,10000]{1,0:T(8,128)} %param, f32[500,10000]{1,0:T(8,128)} %param.1), kind=kLoop, calls=%fused_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"16\",\"40\"],\"input_window_bounds\":[],\"estimated_cycles\":\"142672\",\"iteration_bounds\":[\"4\",\"2\"]},\"scoped_memory_configs\":[]}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2D mesh example\n"
      ],
      "metadata": {
        "id": "_y63HyZIeswz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,4)))"
      ],
      "metadata": {
        "id": "UiCNiPlixOW8"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = jax.device_put(v1s, sharding)\n",
        "v2sp = jax.device_put(v2s, sharding)"
      ],
      "metadata": {
        "id": "_Ts2Yx3rxWui"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "PT7eapRUxiXZ",
        "outputId": "08595b8c-c364-435c-dc22-37ef67130780"
      },
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌─────────────┬─────────────┬─────────────┬─────────────┐\n",
              "│             │             │             │             │\n",
              "│    TPU \u001b[1;36m0\u001b[0m    │    TPU \u001b[1;36m1\u001b[0m    │    TPU \u001b[1;36m2\u001b[0m    │    TPU \u001b[1;36m3\u001b[0m    │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "├─────────────┼─────────────┼─────────────┼─────────────┤\n",
              "│             │             │             │             │\n",
              "│    TPU \u001b[1;36m6\u001b[0m    │    TPU \u001b[1;36m7\u001b[0m    │    TPU \u001b[1;36m4\u001b[0m    │    TPU \u001b[1;36m5\u001b[0m    │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "└─────────────┴─────────────┴─────────────┴─────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌─────────────┬─────────────┬─────────────┬─────────────┐\n",
              "│             │             │             │             │\n",
              "│    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>    │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "├─────────────┼─────────────┼─────────────┼─────────────┤\n",
              "│             │             │             │             │\n",
              "│    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>    │    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>    │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "└─────────────┴─────────────┴─────────────┴─────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)\n",
        "d.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u8ZPzKbfyGE6",
        "outputId": "5def7ebf-2bde-4da3-d1f1-b51ca5d96cf2"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "pflrYiYLyMPQ",
        "outputId": "facdaf00-05a9-4653-fde7-6bae8bbb0a1b"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m│TPU \u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jax.vmap(dot)(v1sp, v2sp).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fiuoqoK2yRhB",
        "outputId": "eccdbb82-91c4-45e5-a7ac-85a922e2f717"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.61 ms ± 39.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jax.vmap(dot)(v1s, v2s).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c_zOa1uyyR3r",
        "outputId": "76f2cf0c-2c63-4275-d3c6-7cf51cb906de"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.32 ms ± 15.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Looking at HLO (note the all-reduce operation)"
      ],
      "metadata": {
        "id": "1XOdpG2MA0Wv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1sp, v2sp).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fGkNyNn52s8T",
        "outputId": "49ac0346-fd7e-41b1-b188-019a66146523"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[2000,2500]{1,0:T(8,128)}, f32[2000,2500]{1,0:T(8,128)})->f32[2000]{0:T(1024)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(256)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(256)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(256)} add(f32[]{:T(256)} %scalar_lhs, f32[]{:T(256)} %scalar_rhs)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.2: f32[2000,2500], param_1.2: f32[2000,2500]) -> f32[2000] {\n",
            "  %param_0.2 = f32[2000,2500]{1,0:T(8,128)} parameter(0)\n",
            "  %param_1.2 = f32[2000,2500]{1,0:T(8,128)} parameter(1)\n",
            "  %multiply.2 = f32[2000,2500]{1,0:T(8,128)} multiply(f32[2000,2500]{1,0:T(8,128)} %param_0.2, f32[2000,2500]{1,0:T(8,128)} %param_1.2)\n",
            "  %constant.2 = f32[]{:T(256)} constant(0)\n",
            "  ROOT %reduce.2 = f32[2000]{0:T(1024)} reduce(f32[2000,2500]{1,0:T(8,128)} %multiply.2, f32[]{:T(256)} %constant.2), dimensions={1}, to_apply=%scalar_add_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[2000,2500], param.1: f32[2000,2500]) -> f32[2000] {\n",
            "  %param.1 = f32[2000,2500]{1,0:T(8,128)} parameter(1), sharding={devices=[2,4]0,1,2,3,4,5,6,7}\n",
            "  %param = f32[2000,2500]{1,0:T(8,128)} parameter(0), sharding={devices=[2,4]0,1,2,3,4,5,6,7}\n",
            "  %fusion = f32[2000]{0:T(1024)} fusion(f32[2000,2500]{1,0:T(8,128)} %param, f32[2000,2500]{1,0:T(8,128)} %param.1), kind=kLoop, calls=%fused_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"32\",\"20\"],\"input_window_bounds\":[],\"estimated_cycles\":\"142736\",\"iteration_bounds\":[\"8\",\"1\"]},\"scoped_memory_configs\":[]}\n",
            "  ROOT %all-reduce = f32[2000]{0:T(1024)} all-reduce(f32[2000]{0:T(1024)} %fusion), channel_id=1, replica_groups={{0,1,2,3},{4,5,6,7}}, use_global_device_ids=true, to_apply=%scalar_add_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"barrier_config\":{\"barrier_type\":\"CUSTOM\",\"id\":\"0\"},\"scoped_memory_configs\":[{\"memory_space\":\"0\",\"offset\":\"0\",\"size\":\"67108864\"}],\"collective_algorithm_config\":{\"emitter\":\"SinglePhaseRingSumEmitter\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Using replication"
      ],
      "metadata": {
        "id": "516tLncUA2Z5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,4)))\n",
        "\n",
        "v1sp = jax.device_put(v1s, sharding.replicate(axis=1))\n",
        "\n",
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "ooOHrGv64TZF",
        "outputId": "539aa17a-df53-4186-9e00-39c5a62d53d0"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                        TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m                         │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                                                            │\n",
              "│                        TPU \u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m                         │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                        TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                         │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                                                            │\n",
              "│                        TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                         │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "A = random.normal(rng_key, shape=(10000,2000))\n",
        "B = random.normal(rng_key, shape=(2000,5000))"
      ],
      "metadata": {
        "id": "Y1wPyLI1C6P3"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "Ad = jax.device_put(A, sharding.replicate(1))\n",
        "Bd = jax.device_put(B, sharding.replicate(0))"
      ],
      "metadata": {
        "id": "qFct_kR6DFho"
      },
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(Ad)\n",
        "jax.debug.visualize_array_sharding(Bd)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "OYhgu1U4DVYg",
        "outputId": "68ffb6bb-d1f1-4d27-be31-f2a58090ab61"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┐\n",
              "│           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m│\n",
              "│           │\n",
              "│           │\n",
              "├───────────┤\n",
              "│           │\n",
              "│TPU \u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │\n",
              "│           │\n",
              "└───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┐\n",
              "│           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>│\n",
              "│           │\n",
              "│           │\n",
              "├───────────┤\n",
              "│           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │\n",
              "│           │\n",
              "└───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌─────────────┬─────────────┬─────────────┬─────────────┐\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│   TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m6\u001b[0m   │   TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m7\u001b[0m   │   TPU \u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m   │   TPU \u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m   │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "└─────────────┴─────────────┴─────────────┴─────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌─────────────┬─────────────┬─────────────┬─────────────┐\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│   TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>   │   TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>   │   TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>   │   TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>   │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "│             │             │             │             │\n",
              "└─────────────┴─────────────┴─────────────┴─────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Cd = jnp.dot(Ad, Bd)"
      ],
      "metadata": {
        "id": "ld4UAh1QDOxU"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(Cd)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "QdzAVah9DXyA",
        "outputId": "0b002581-f0dc-4da9-bb46-3a95d3bb2853"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┬───────┬───────┬───────┐\n",
              "│       │       │       │       │\n",
              "│ TPU \u001b[1;36m0\u001b[0m │ TPU \u001b[1;36m1\u001b[0m │ TPU \u001b[1;36m2\u001b[0m │ TPU \u001b[1;36m3\u001b[0m │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "├───────┼───────┼───────┼───────┤\n",
              "│       │       │       │       │\n",
              "│ TPU \u001b[1;36m6\u001b[0m │ TPU \u001b[1;36m7\u001b[0m │ TPU \u001b[1;36m4\u001b[0m │ TPU \u001b[1;36m5\u001b[0m │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "└───────┴───────┴───────┴───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
              "│       │       │       │       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span> │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "├───────┼───────┼───────┼───────┤\n",
              "│       │       │       │       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span> │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "└───────┴───────┴───────┴───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "C = A @ B"
      ],
      "metadata": {
        "id": "_W-UMoSkDZA2"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.array_equal(C,Cd)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MNHr5cE8DcS2",
        "outputId": "4c7571a8-8dcb-4b98-f302-bf756df8e40f"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(C)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "RI9dX5LlE7sh",
        "outputId": "d1878349-77f6-449c-9ff2-36ce6c3c98f9"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU \u001b[1;36m0\u001b[0m   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "C.shape, Cd.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DF3Gv28IVO6V",
        "outputId": "7c51e884-2ba2-46af-ec9e-e684e97c93d5"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((10000, 5000), (10000, 5000))"
            ]
          },
          "metadata": {},
          "execution_count": 46
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "C[12,3], Cd[12,3]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RtVhE7KKVeji",
        "outputId": "e1222f9f-09e7-48f6-a036-65d22cfb62e2"
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Array(43.027637, dtype=float32), Array(43.027637, dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 47
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Ca = jnp.dot(A, B)"
      ],
      "metadata": {
        "id": "QxgSv_ZSVquX"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.array_equal(C,Ca)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xO_sxUFrVtpN",
        "outputId": "91b814fd-dd30-43cf-c98d-e603ffa3feb9"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(Ca)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "wQvIChAAV0f_",
        "outputId": "5ec8e7f5-04d3-444c-ba7e-893abb4fef4e"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU \u001b[1;36m0\u001b[0m   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = (Cd - C)"
      ],
      "metadata": {
        "id": "UdXSxOhHV7wl"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.max(d), jnp.sum(d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I-uXR9-EV_A1",
        "outputId": "5f3937e4-cd16-4b27-cfcf-d5232a441443"
      },
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Array(0., dtype=float32), Array(0., dtype=float32))"
            ]
          },
          "metadata": {},
          "execution_count": 52
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jnp.dot(Ad, Bd).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6RCR-BcuWtvi",
        "outputId": "31f53d02-e4ce-4db3-9038-e994292b1634"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.12 ms ± 58.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit jnp.dot(A, B).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GSS3YK4IW21p",
        "outputId": "0510c115-6b9f-4673-8550-de7d59acb2c7"
      },
      "execution_count": 54,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10.4 ms ± 10.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit (A@B).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fF_EaKAGW5Nk",
        "outputId": "290105b2-8a8d-4a16-ab0e-7f4965b91afe"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10.4 ms ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "modules = jnp.dot.lower(A, B).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Nrq6pY5fH3Id",
        "outputId": "af8db4e2-6f36-4184-a9e6-d6885c4a76dd"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[10000,2000]{0,1:T(8,128)}, f32[2000,5000]{1,0:T(8,128)})->f32[10000,5000]{0,1:T(8,128)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%bitcast_fusion (bf16input: f32[10000,2000]) -> f32[10000,2000] {\n",
            "  %bf16input = f32[10000,2000]{0,1:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast = f32[10000,2000]{0,1:T(8,128)} bitcast(f32[10000,2000]{0,1:T(8,128)} %bf16input)\n",
            "}\n",
            "\n",
            "%bitcast_fusion.1 (bf16input.1: f32[2000,5000]) -> f32[2000,5000] {\n",
            "  %bf16input.1 = f32[2000,5000]{1,0:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast.1 = f32[2000,5000]{1,0:T(8,128)} bitcast(f32[2000,5000]{1,0:T(8,128)} %bf16input.1)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0: f32[10000,2000], param_1: f32[2000,5000]) -> f32[10000,5000] {\n",
            "  %param_0 = f32[10000,2000]{0,1:T(8,128)} parameter(0)\n",
            "  %fusion.1 = f32[10000,2000]{0,1:T(8,128)} fusion(f32[10000,2000]{0,1:T(8,128)} %param_0), kind=kLoop, calls=%bitcast_fusion\n",
            "  %param_1 = f32[2000,5000]{1,0:T(8,128)} parameter(1)\n",
            "  %fusion.2 = f32[2000,5000]{1,0:T(8,128)} fusion(f32[2000,5000]{1,0:T(8,128)} %param_1), kind=kLoop, calls=%bitcast_fusion.1\n",
            "  ROOT %convolution.1 = f32[10000,5000]{0,1:T(8,128)} convolution(f32[10000,2000]{0,1:T(8,128)} %fusion.1, f32[2000,5000]{1,0:T(8,128)} %fusion.2), dim_labels=bf_io->bf, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1507742915.py\" source_line=1}\n",
            "}\n",
            "\n",
            "ENTRY %main.4 (Arg_0.1: f32[10000,2000], Arg_1.2: f32[2000,5000]) -> f32[10000,5000] {\n",
            "  %Arg_1.2 = f32[2000,5000]{1,0:T(8,128)} parameter(1), sharding={replicated}\n",
            "  %Arg_0.1 = f32[10000,2000]{0,1:T(8,128)} parameter(0), sharding={replicated}\n",
            "  ROOT %fusion = f32[10000,5000]{0,1:T(8,128)} fusion(f32[10000,2000]{0,1:T(8,128)} %Arg_0.1, f32[2000,5000]{1,0:T(8,128)} %Arg_1.2), kind=kOutput, calls=%fused_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1507742915.py\" source_line=1}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"250\",\"4\"],\"output_window_bounds\":[\"64\",\"1\"],\"input_window_bounds\":[\"250\",\"1\"],\"estimated_cycles\":\"6747958\",\"iteration_bounds\":[\"10\",\"79\",\"1\"]},\"scoped_memory_configs\":[]}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "modules = jnp.dot.lower(Ad, Bd).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P1auVbzQHxAG",
        "outputId": "54477235-7c3a-497b-c9db-f67293e20a14"
      },
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[5000,2000]{1,0:T(8,128)}, f32[2000,1250]{1,0:T(8,128)})->f32[5000,1250]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%bitcast_fusion (bf16input: f32[5000,2000]) -> f32[5000,2000] {\n",
            "  %bf16input = f32[5000,2000]{1,0:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast = f32[5000,2000]{1,0:T(8,128)} bitcast(f32[5000,2000]{1,0:T(8,128)} %bf16input)\n",
            "}\n",
            "\n",
            "%bitcast_fusion.1 (bf16input.1: f32[2000,1250]) -> f32[2000,1250] {\n",
            "  %bf16input.1 = f32[2000,1250]{1,0:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast.1 = f32[2000,1250]{1,0:T(8,128)} bitcast(f32[2000,1250]{1,0:T(8,128)} %bf16input.1)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0: f32[5000,2000], param_1: f32[2000,1250]) -> f32[5000,1250] {\n",
            "  %param_0 = f32[5000,2000]{1,0:T(8,128)} parameter(0)\n",
            "  %fusion.1 = f32[5000,2000]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)} %param_0), kind=kLoop, calls=%bitcast_fusion\n",
            "  %param_1 = f32[2000,1250]{1,0:T(8,128)} parameter(1)\n",
            "  %fusion.2 = f32[2000,1250]{1,0:T(8,128)} fusion(f32[2000,1250]{1,0:T(8,128)} %param_1), kind=kLoop, calls=%bitcast_fusion.1\n",
            "  ROOT %convolution.1 = f32[5000,1250]{1,0:T(8,128)} convolution(f32[5000,2000]{1,0:T(8,128)} %fusion.1, f32[2000,1250]{1,0:T(8,128)} %fusion.2), dim_labels=bf_io->bf, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1507742915.py\" source_line=1}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[5000,2000], param.1: f32[2000,1250]) -> f32[5000,1250] {\n",
            "  %param.1 = f32[2000,1250]{1,0:T(8,128)} parameter(1), sharding={devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}\n",
            "  %param = f32[5000,2000]{1,0:T(8,128)} parameter(0), sharding={devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}\n",
            "  ROOT %fusion = f32[5000,1250]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)} %param, f32[2000,1250]{1,0:T(8,128)} %param.1), kind=kOutput, calls=%fused_computation, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1507742915.py\" source_line=1}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"64\",\"5\"],\"output_window_bounds\":[\"125\",\"5\"],\"input_window_bounds\":[\"125\",\"4\"],\"estimated_cycles\":\"837631\",\"iteration_bounds\":[\"2\",\"5\",\"4\"]},\"scoped_memory_configs\":[]}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Using sharding constraints"
      ],
      "metadata": {
        "id": "USf1avNxJ8Na"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import jit\n",
        "from functools import partial\n",
        "\n",
        "@partial(jax.jit, static_argnums=2)\n",
        "def distributed_mul(a, b, sharding):\n",
        "  ad = jax.lax.with_sharding_constraint(a, sharding.replicate(1))\n",
        "  bd = jax.lax.with_sharding_constraint(b, sharding.replicate(0))\n",
        "  return jnp.dot(ad, bd)"
      ],
      "metadata": {
        "id": "GquL4IMFJ-rd"
      },
      "execution_count": 59,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,4)))"
      ],
      "metadata": {
        "id": "7LEVlx4akw5c"
      },
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(A)\n",
        "jax.debug.visualize_array_sharding(B)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "G1ThCrDikrjL",
        "outputId": "70966bbc-d5e2-45ba-9e51-9f889a591850"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┐\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│ TPU \u001b[1;36m0\u001b[0m │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "└───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┐\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "└───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "│                                                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = distributed_mul(A, B, sharding)"
      ],
      "metadata": {
        "id": "HiJOk9HEk7kT"
      },
      "execution_count": 62,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "knUitJvSk_1Z",
        "outputId": "37c86573-0a41-458e-d94f-60bab8e578c8"
      },
      "execution_count": 63,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┬───────┬───────┬───────┐\n",
              "│       │       │       │       │\n",
              "│ TPU \u001b[1;36m0\u001b[0m │ TPU \u001b[1;36m1\u001b[0m │ TPU \u001b[1;36m2\u001b[0m │ TPU \u001b[1;36m3\u001b[0m │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "├───────┼───────┼───────┼───────┤\n",
              "│       │       │       │       │\n",
              "│ TPU \u001b[1;36m6\u001b[0m │ TPU \u001b[1;36m7\u001b[0m │ TPU \u001b[1;36m4\u001b[0m │ TPU \u001b[1;36m5\u001b[0m │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "└───────┴───────┴───────┴───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┐\n",
              "│       │       │       │       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span> │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "├───────┼───────┼───────┼───────┤\n",
              "│       │       │       │       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span> │\n",
              "│       │       │       │       │\n",
              "│       │       │       │       │\n",
              "└───────┴───────┴───────┴───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fkhWKJZBnPC0",
        "outputId": "75b0090d-d5bc-4bff-aec9-945705258aba"
      },
      "execution_count": 64,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000, 5000)"
            ]
          },
          "metadata": {},
          "execution_count": 64
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "@jit\n",
        "def nondistributed_mul(a, b):\n",
        "  return jnp.dot(a, b)"
      ],
      "metadata": {
        "id": "_ArK0PqwlgOj"
      },
      "execution_count": 65,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dn = nondistributed_mul(A, B)"
      ],
      "metadata": {
        "id": "EGSWUis-liuu"
      },
      "execution_count": 66,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(dn)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "PrAZODtAllOl",
        "outputId": "449d3c52-b782-47ee-91e4-caf6d391ad6e"
      },
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU \u001b[1;36m0\u001b[0m   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────┐\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│  TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>   │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "│          │\n",
              "└──────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit distributed_mul(A, B, sharding).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JIo3eJqklrmG",
        "outputId": "f1bc06a7-e3d3-48a9-cdeb-d665340f8318"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "21.2 ms ± 101 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit nondistributed_mul(A, B).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wmyvZXW2luce",
        "outputId": "bfb69e03-8251-4dd1-951d-301c270ca23e"
      },
      "execution_count": 69,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10.4 ms ± 9.82 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "modules = distributed_mul.lower(A, B, sharding).compile().compiler_ir()\n",
        "for hlo in modules:\n",
        "  print(hlo.to_string())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YXUu0WcaHf0H",
        "outputId": "f2b131d0-b416-4450-a61a-b5e7c9b05b95"
      },
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HloModule jit_distributed_mul, is_scheduled=true, entry_computation_layout={(f32[10000,2000]{0,1:T(8,128)}, f32[2000,5000]{1,0:T(8,128)})->f32[5000,1250]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%fused_computation.1 (param_0.2: f32[10000,2000], param_1.2: s32[]) -> bf16[5000,2000] {\n",
            "  %param_0.2 = f32[10000,2000]{0,1:T(8,128)} parameter(0)\n",
            "  %param_1.2 = s32[]{:T(256)} parameter(1)\n",
            "  %constant.23 = s32[]{:T(256)} constant(0), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  ROOT %dynamic-slice.8 = bf16[5000,2000]{0,1:T(8,128)(2,1)} dynamic-slice(f32[10000,2000]{0,1:T(8,128)} %param_0.2, s32[]{:T(256)} %param_1.2, s32[]{:T(256)} %constant.23), dynamic_slice_sizes={5000,2000}, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}, backend_config={\"flag_configs\":[],\"scoped_memory_configs\":[],\"indices_config\":{\"index_known_bits\":[{\"zeroes\":\"7\",\"ones\":\"0\",\"bitwidth\":\"32\"},{\"zeroes\":\"4294967295\",\"ones\":\"0\",\"bitwidth\":\"32\"}]}}\n",
            "}\n",
            "\n",
            "%fused_computation.2 (param_0.4: f32[2000,5000], param_1.5: s32[]) -> bf16[2000,1250] {\n",
            "  %param_0.4 = f32[2000,5000]{1,0:T(8,128)} parameter(0)\n",
            "  %constant.26 = s32[]{:T(256)} constant(0), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %param_1.5 = s32[]{:T(256)} parameter(1)\n",
            "  ROOT %dynamic-slice.10 = bf16[2000,1250]{1,0:T(8,128)(2,1)} dynamic-slice(f32[2000,5000]{1,0:T(8,128)} %param_0.4, s32[]{:T(256)} %constant.26, s32[]{:T(256)} %param_1.5), dynamic_slice_sizes={2000,1250}, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}, backend_config={\"flag_configs\":[],\"scoped_memory_configs\":[],\"indices_config\":{\"index_known_bits\":[{\"zeroes\":\"4294967295\",\"ones\":\"0\",\"bitwidth\":\"32\"},{\"zeroes\":\"1\",\"ones\":\"0\",\"bitwidth\":\"32\"}]}}\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.3: f32[10000,2000], param_1.3: s32[], param_2.4: f32[2000,5000], param_3.2: s32[]) -> f32[5000,1250] {\n",
            "  %param_0.3 = f32[10000,2000]{0,1:T(8,128)} parameter(0)\n",
            "  %param_1.3 = s32[]{:T(256)} parameter(1)\n",
            "  %fusion.1 = bf16[5000,2000]{0,1:T(8,128)(2,1)} fusion(f32[10000,2000]{0,1:T(8,128)} %param_0.3, s32[]{:T(256)} %param_1.3), kind=kLoop, calls=%fused_computation.1, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %param_2.4 = f32[2000,5000]{1,0:T(8,128)} parameter(2)\n",
            "  %param_3.2 = s32[]{:T(256)} parameter(3)\n",
            "  %fusion.2 = bf16[2000,1250]{1,0:T(8,128)(2,1)} fusion(f32[2000,5000]{1,0:T(8,128)} %param_2.4, s32[]{:T(256)} %param_3.2), kind=kLoop, calls=%fused_computation.2, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  ROOT %convolution.1 = f32[5000,1250]{1,0:T(8,128)} convolution(bf16[5000,2000]{0,1:T(8,128)(2,1)} %fusion.1, bf16[2000,1250]{1,0:T(8,128)(2,1)} %fusion.2), dim_labels=bf_io->bf, metadata={op_name=\"jit(distributed_mul)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=8}\n",
            "}\n",
            "\n",
            "ENTRY %main.6_spmd (param: f32[10000,2000], param.1: f32[2000,5000]) -> f32[5000,1250] {\n",
            "  %constant.18 = s32[1]{0:T(256)} constant({1250})\n",
            "  %constant.17 = s32[1]{0:T(256)} constant({5000})\n",
            "  %constant.16 = u32[1]{0:T(256)} constant({3})\n",
            "  %constant.14 = u32[1]{0:T(256)} constant({1})\n",
            "  %constant.13 = u32[1]{0:T(256)} constant({0})\n",
            "  %constant.3 = u32[8]{0:T(256)} constant({0, 1, 2, 3, 0, 1, 2, 3}), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  %constant = u32[8]{0:T(256)} constant({0, 0, 0, 0, 1, 1, 1, 1}), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %param.1 = f32[2000,5000]{1,0:T(8,128)} parameter(1), sharding={replicated}\n",
            "  %param = f32[10000,2000]{0,1:T(8,128)} parameter(0), sharding={replicated}\n",
            "  %partition-id = u32[]{:T(256)} partition-id()\n",
            "  %dynamic-slice.3 = u32[1]{0:T(256)} dynamic-slice(u32[8]{0:T(256)} %constant.3, u32[]{:T(256)} %partition-id), dynamic_slice_sizes={1}, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"1\"],\"input_window_bounds\":[],\"estimated_cycles\":\"183\",\"iteration_bounds\":[\"1\"]},\"scoped_memory_configs\":[],\"indices_config\":{\"index_known_bits\":[{\"zeroes\":\"0\",\"ones\":\"0\",\"bitwidth\":\"32\"}]}}\n",
            "  %clamp.3 = u32[1]{0:T(256)S(6)} clamp(u32[1]{0:T(256)} %constant.13, u32[1]{0:T(256)} %dynamic-slice.3, u32[1]{0:T(256)} %constant.16), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  %convert.3 = s32[1]{0:T(256)S(6)} convert(u32[1]{0:T(256)S(6)} %clamp.3), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  %multiply.5 = s32[1]{0:T(256)} multiply(s32[1]{0:T(256)S(6)} %convert.3, s32[1]{0:T(256)} %constant.18), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  %bitcast.1 = s32[]{:T(256)} bitcast(s32[1]{0:T(256)} %multiply.5), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[1,4,2]0,4,1,5,2,6,3,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=7}\n",
            "  %dynamic-slice = u32[1]{0:T(256)} dynamic-slice(u32[8]{0:T(256)} %constant, u32[]{:T(256)} %partition-id), dynamic_slice_sizes={1}, metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"1\"],\"input_window_bounds\":[],\"estimated_cycles\":\"183\",\"iteration_bounds\":[\"1\"]},\"scoped_memory_configs\":[],\"indices_config\":{\"index_known_bits\":[{\"zeroes\":\"0\",\"ones\":\"0\",\"bitwidth\":\"32\"}]}}\n",
            "  %clamp.2 = u32[1]{0:T(256)S(6)} clamp(u32[1]{0:T(256)} %constant.13, u32[1]{0:T(256)} %dynamic-slice, u32[1]{0:T(256)} %constant.14), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %convert.2 = s32[1]{0:T(256)S(6)} convert(u32[1]{0:T(256)S(6)} %clamp.2), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %multiply.4 = s32[1]{0:T(256)} multiply(s32[1]{0:T(256)S(6)} %convert.2, s32[1]{0:T(256)} %constant.17), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  %bitcast = s32[]{:T(256)} bitcast(s32[1]{0:T(256)} %multiply.4), metadata={op_name=\"jit(distributed_mul)/jit(main)/sharding_constraint[sharding=GSPMDSharding({devices=[2,1,4]0,1,2,3,4,5,6,7 last_tile_dim_replicate}) resource_env=ResourceEnv(Mesh(device_ids=[], axis_names=()), ()) unconstrained_dims={}]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=6}\n",
            "  ROOT %fusion = f32[5000,1250]{1,0:T(8,128)} fusion(f32[10000,2000]{0,1:T(8,128)} %param, s32[]{:T(256)} %bitcast, f32[2000,5000]{1,0:T(8,128)} %param.1, s32[]{:T(256)} %bitcast.1), kind=kOutput, calls=%fused_computation, metadata={op_name=\"jit(distributed_mul)/jit(main)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=None]\" source_file=\"/tmp/ipykernel_10444/1742017423.py\" source_line=8}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"63\",\"10\"],\"output_window_bounds\":[\"64\",\"10\"],\"input_window_bounds\":[\"63\",\"4\"],\"estimated_cycles\":\"856748\",\"iteration_bounds\":[\"1\",\"10\",\"4\"]},\"scoped_memory_configs\":[]}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Named sharding"
      ],
      "metadata": {
        "id": "rPIpRTgO9fDa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.sharding import Mesh\n",
        "from jax.sharding import PartitionSpec as P\n",
        "from jax.sharding import NamedSharding"
      ],
      "metadata": {
        "id": "BbfZ4Zto9gLY"
      },
      "execution_count": 71,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "mesh = Mesh(mesh_utils.create_device_mesh((4,2)), axis_names=('batch', 'features'))\n",
        "sharding = NamedSharding(mesh, P('batch', 'features'))"
      ],
      "metadata": {
        "id": "Pz42EU019j99"
      },
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = jax.device_put(v1s, sharding)\n",
        "v2sp = jax.device_put(v2s, sharding)\n",
        "\n",
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "iu1O52k5-caD",
        "outputId": "4ad2e8b0-d672-4489-c6c6-14d9a3557d5c"
      },
      "execution_count": 73,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌─────────────────────────────┬─────────────────────────────┐\n",
              "│            TPU \u001b[1;36m0\u001b[0m            │            TPU \u001b[1;36m1\u001b[0m            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU \u001b[1;36m2\u001b[0m            │            TPU \u001b[1;36m3\u001b[0m            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU \u001b[1;36m6\u001b[0m            │            TPU \u001b[1;36m7\u001b[0m            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU \u001b[1;36m4\u001b[0m            │            TPU \u001b[1;36m5\u001b[0m            │\n",
              "└─────────────────────────────┴─────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌─────────────────────────────┬─────────────────────────────┐\n",
              "│            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>            │            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>            │            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>            │            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>            │\n",
              "├─────────────────────────────┼─────────────────────────────┤\n",
              "│            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>            │            TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>            │\n",
              "└─────────────────────────────┴─────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)\n",
        "d.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "18uhGU7pAgui",
        "outputId": "aed16b54-a4ed-4aa2-f057-30ce3ea74a78"
      },
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "metadata": {},
          "execution_count": 74
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Device placement policy and errors"
      ],
      "metadata": {
        "id": "RBjyztjhKLxx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Different devices:"
      ],
      "metadata": {
        "id": "zCpzsk5vbNne"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding_a = PositionalSharding(np.array(jax.devices()[:4]).reshape(4,1))\n",
        "sharding_b = PositionalSharding(np.array(jax.devices()[4:]).reshape(4,1))"
      ],
      "metadata": {
        "id": "qqIQJ2ejKJw9"
      },
      "execution_count": 75,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sharding_a"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JAbiGQEUaE4d",
        "outputId": "45fa969a-47a3-401f-e355-59ba4098a4c6"
      },
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 0}]\n",
              "                    [{TPU 1}]\n",
              "                    [{TPU 2}]\n",
              "                    [{TPU 3}]])"
            ]
          },
          "metadata": {},
          "execution_count": 76
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sharding_b"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7lmmW2UkaFqY",
        "outputId": "9b67c75f-c363-4a2e-de3e-31450276b59e"
      },
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 4}]\n",
              "                    [{TPU 5}]\n",
              "                    [{TPU 6}]\n",
              "                    [{TPU 7}]])"
            ]
          },
          "metadata": {},
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = jax.device_put(v1s, sharding_a)\n",
        "v2sp = jax.device_put(v2s, sharding_b)"
      ],
      "metadata": {
        "id": "VQ-IkTSmaKst"
      },
      "execution_count": 78,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "cLk4ddMBboyI",
        "outputId": "01b27140-aee2-41fa-cedc-02b9d4f61bd3"
      },
      "execution_count": 79,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m1\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m2\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m3\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "88sPq_RZbpnY",
        "outputId": "5504d57a-8218-4bdf-8f16-e89a03e9b126"
      },
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m4\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m5\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m6\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m7\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 453
        },
        "id": "WYBV-u87aNs1",
        "outputId": "352a452e-244e-49a2-ec9e-28e2b2dce746"
      },
      "execution_count": 81,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1, 2, 3] on platform TPU and ARG_SHARDING with device ids [4, 5, 6, 7] on platform TPU",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[81], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m d \u001b[38;5;241m=\u001b[39m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdot\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1sp\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2sp\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 3 frame]\u001b[0m\n",
            "Cell \u001b[0;32mIn[9], line 2\u001b[0m, in \u001b[0;36mdot\u001b[0;34m(v1, v2)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdot\u001b[39m(v1, v2):\n\u001b[0;32m----> 2\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvdot\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 2 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/pjit.py:171\u001b[0m, in \u001b[0;36m_python_pjit_helper\u001b[0;34m(fun, infer_params_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m    168\u001b[0m   fun_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(fun, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__qualname__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28mgetattr\u001b[39m(fun, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28mstr\u001b[39m(fun)))\n\u001b[1;32m    169\u001b[0m   msg \u001b[38;5;241m=\u001b[39m _device_assignment_mismatch_error(\n\u001b[1;32m    170\u001b[0m       fun_name, fails, args_flat, api_name, arg_names)\n\u001b[0;32m--> 171\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    172\u001b[0m outs \u001b[38;5;241m=\u001b[39m tree_unflatten(out_tree, out_flat)\n\u001b[1;32m    173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m outs, out_flat, out_tree, args_flat, params[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mjaxpr\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
            "\u001b[0;31mValueError\u001b[0m: Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1, 2, 3] on platform TPU and ARG_SHARDING with device ids [4, 5, 6, 7] on platform TPU"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Different order:"
      ],
      "metadata": {
        "id": "0C5PKKelbP8n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding_a = PositionalSharding(np.array(jax.devices()).reshape(8,1))\n",
        "sharding_b = PositionalSharding(np.array(jax.devices()[::-1]).reshape(8,1))"
      ],
      "metadata": {
        "id": "XV6jtqLSdKqC"
      },
      "execution_count": 82,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sharding_a, sharding_b"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-2ly9uCNdODU",
        "outputId": "23d8b552-2453-4235-ebb3-a8bbf80e6630"
      },
      "execution_count": 83,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(PositionalSharding([[{TPU 0}]\n",
              "                     [{TPU 1}]\n",
              "                     [{TPU 2}]\n",
              "                     [{TPU 3}]\n",
              "                     [{TPU 4}]\n",
              "                     [{TPU 5}]\n",
              "                     [{TPU 6}]\n",
              "                     [{TPU 7}]]),\n",
              " PositionalSharding([[{TPU 7}]\n",
              "                     [{TPU 6}]\n",
              "                     [{TPU 5}]\n",
              "                     [{TPU 4}]\n",
              "                     [{TPU 3}]\n",
              "                     [{TPU 2}]\n",
              "                     [{TPU 1}]\n",
              "                     [{TPU 0}]]))"
            ]
          },
          "metadata": {},
          "execution_count": 83
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = jax.device_put(v1s, sharding_a)\n",
        "v2sp = jax.device_put(v2s, sharding_b)"
      ],
      "metadata": {
        "id": "jW8JqD73dW0u"
      },
      "execution_count": 84,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "5MM-Gd4fdXL4",
        "outputId": "6b1d87be-acb9-4db3-fa96-4e035dc77a0f"
      },
      "execution_count": 85,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m1\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m2\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m3\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m4\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m5\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m6\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m7\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "Vzm3d9OWdZ7b",
        "outputId": "4b31a0c5-a221-498a-f6cf-81704f447649"
      },
      "execution_count": 86,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU \u001b[1;36m7\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m6\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m5\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m4\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m3\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m2\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m1\u001b[0m                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU \u001b[1;36m0\u001b[0m                            │\n",
              "└────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌────────────────────────────────────────────────────────────┐\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>                            │\n",
              "├────────────────────────────────────────────────────────────┤\n",
              "│                           TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                            │\n",
              "└────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 453
        },
        "id": "JsIjQdCtdeIg",
        "outputId": "e51aadfd-e770-4a3d-81c8-f854f66ba604"
      },
      "execution_count": 87,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1, 2, 3, 4, 5, 6, 7] on platform TPU and ARG_SHARDING with device ids [7, 6, 5, 4, 3, 2, 1, 0] on platform TPU",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[87], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m d \u001b[38;5;241m=\u001b[39m \u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdot\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1sp\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2sp\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 3 frame]\u001b[0m\n",
            "Cell \u001b[0;32mIn[9], line 2\u001b[0m, in \u001b[0;36mdot\u001b[0;34m(v1, v2)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdot\u001b[39m(v1, v2):\n\u001b[0;32m----> 2\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mjnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvdot\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv1\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mv2\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 2 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/pjit.py:171\u001b[0m, in \u001b[0;36m_python_pjit_helper\u001b[0;34m(fun, infer_params_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m    168\u001b[0m   fun_name \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(fun, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__qualname__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28mgetattr\u001b[39m(fun, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28mstr\u001b[39m(fun)))\n\u001b[1;32m    169\u001b[0m   msg \u001b[38;5;241m=\u001b[39m _device_assignment_mismatch_error(\n\u001b[1;32m    170\u001b[0m       fun_name, fails, args_flat, api_name, arg_names)\n\u001b[0;32m--> 171\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    172\u001b[0m outs \u001b[38;5;241m=\u001b[39m tree_unflatten(out_tree, out_flat)\n\u001b[1;32m    173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m outs, out_flat, out_tree, args_flat, params[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mjaxpr\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
            "\u001b[0;31mValueError\u001b[0m: Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1, 2, 3, 4, 5, 6, 7] on platform TPU and ARG_SHARDING with device ids [7, 6, 5, 4, 3, 2, 1, 0] on platform TPU"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "d = jax.vmap(dot)(v1sp, v2s)"
      ],
      "metadata": {
        "id": "_Z7Mlk2KfPW3"
      },
      "execution_count": 88,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "FdcPLXpPfQrc",
        "outputId": "74277cd0-92a0-4c6f-d3c7-acdc5804bb2f"
      },
      "execution_count": 89,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐\n",
              "│ TPU \u001b[1;36m0\u001b[0m │ TPU \u001b[1;36m1\u001b[0m │ TPU \u001b[1;36m2\u001b[0m │ TPU \u001b[1;36m3\u001b[0m │ TPU \u001b[1;36m4\u001b[0m │ TPU \u001b[1;36m5\u001b[0m │ TPU \u001b[1;36m6\u001b[0m │ TPU \u001b[1;36m7\u001b[0m │\n",
              "└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span> │ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span> │\n",
              "└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "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": "cd9a364e-17d2-499e-e01d-67d100c1d7ed"
      },
      "execution_count": 169,
      "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     |████████████████████████████████| 479.6 MB 10 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 12.4 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 6.2 MB/s \n",
            "\u001b[?25hCollecting grpcio<2.0,>=1.24.3\n",
            "  Downloading grpcio-1.60.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.4 MB 12.7 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 13.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 14.4 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 14.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 7.7 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.25.1-cp37-abi3-manylinux2014_x86_64.whl (294 kB)\n",
            "\u001b[K     |████████████████████████████████| 294 kB 6.1 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 2.5 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 7.1 MB/s \n",
            "\u001b[?25hCollecting termcolor>=1.1.0\n",
            "  Downloading termcolor-2.4.0-py3-none-any.whl (7.7 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.16.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (83 kB)\n",
            "\u001b[K     |████████████████████████████████| 83 kB 2.0 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 6.7 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.26.1-py2.py3-none-any.whl (186 kB)\n",
            "\u001b[K     |████████████████████████████████| 186 kB 7.6 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.1-py3-none-any.whl (102 kB)\n",
            "\u001b[K     |████████████████████████████████| 102 kB 5.7 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 7.2 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) (7.0.1)\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.9.0\n",
            "    Uninstalling typing-extensions-4.9.0:\n",
            "      Successfully uninstalled typing-extensions-4.9.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.26.1 google-auth-oauthlib-1.0.0 google-pasta-0.2.0 grpcio-1.60.0 h5py-3.10.0 keras-2.13.1 libclang-16.0.6 markdown-3.5.1 numpy-1.24.3 protobuf-4.25.1 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.4.0 typing-extensions-4.5.0 werkzeug-3.0.1 wrapt-1.16.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow_datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R1PpLK59K12_",
        "outputId": "8189efb7-51ce-4e29-d5ab-ff91d4952a54"
      },
      "execution_count": 170,
      "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 1.6 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 94.6 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 84.1 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 67.4 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.25.1)\n",
            "Requirement already satisfied: psutil in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (5.9.7)\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.4.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.16.0)\n",
            "Requirement already satisfied: importlib-resources; python_version < \"3.9\" in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (6.1.1)\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.62.0-py2.py3-none-any.whl (228 kB)\n",
            "\u001b[K     |████████████████████████████████| 228 kB 75.2 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=dbbd2017f01bac232a9718cc212ea54ee86a03bc81f18387e6ba90498c474414\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.25.1 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.62.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",
      "execution_count": 171,
      "metadata": {
        "id": "BrWQqOqTuz8C",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6e0718e6-b014-4544-d229-8c2890d8c179"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-01-10 11:43:24.979689: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
            "/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.12 file/s]"
          ]
        },
        {
          "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"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\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": 172,
      "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",
        "BATCH_SIZE  = 32 # total 60k samples\n",
        "NUM_DEVICES = jax.device_count()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 173,
      "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": 174,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "seqmmqqx1let",
        "outputId": "1df1ebe1-264f-4518-da9f-4656104d4299"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "235"
            ]
          },
          "metadata": {},
          "execution_count": 174
        }
      ],
      "source": [
        "len(train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TkQT5al_1ove"
      },
      "source": [
        "### Preparing MLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 175,
      "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": 176,
      "metadata": {
        "id": "oRrZN0Xh1qsl"
      },
      "outputs": [],
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "PARAM_SCALE = 0.01"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 177,
      "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": 178,
      "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": 179,
      "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": 180,
      "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",
        "@jit\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": "markdown",
      "metadata": {
        "id": "XFyXVUe62Epe"
      },
      "source": [
        "### Training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 181,
      "metadata": {
        "id": "MjDtmztL1925"
      },
      "outputs": [],
      "source": [
        "@jit\n",
        "def batch_accuracy(params, images, targets):\n",
        "  images = jnp.reshape(images, (len(images), NUM_PIXELS))\n",
        "  predicted_class = jnp.argmax(batched_predict(params, images), axis=1)\n",
        "  return jnp.mean(predicted_class == targets)\n",
        "\n",
        "def accuracy(params, data):\n",
        "  accs = []\n",
        "  for images, targets in data:\n",
        "    accs.append(batch_accuracy(params, images, targets))\n",
        "  return jnp.mean(jnp.array(accs))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 8-way data parallelism"
      ],
      "metadata": {
        "id": "hzF2I-4V3Qjv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(jax.devices()).reshape(8, 1)"
      ],
      "metadata": {
        "id": "JrhlyerwzgZ1"
      },
      "execution_count": 182,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 183,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CGd4pczh2I1h",
        "outputId": "8a870220-484d-4095-ceb2-1f0402a7daa9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0 in 1.91 sec\n",
            "Training set loss 0.6914460062980652\n",
            "Training set accuracy 0.8648548126220703\n",
            "Test set accuracy 0.874804675579071\n",
            "Epoch 1 in 1.00 sec\n",
            "Training set loss 0.6259850859642029\n",
            "Training set accuracy 0.8882258534431458\n",
            "Test set accuracy 0.893847644329071\n",
            "Epoch 2 in 0.99 sec\n",
            "Training set loss 0.6164931058883667\n",
            "Training set accuracy 0.8978224396705627\n",
            "Test set accuracy 0.901074230670929\n",
            "Epoch 3 in 1.00 sec\n",
            "Training set loss 0.6100647449493408\n",
            "Training set accuracy 0.906527042388916\n",
            "Test set accuracy 0.908984363079071\n",
            "Epoch 4 in 0.97 sec\n",
            "Training set loss 0.6044379472732544\n",
            "Training set accuracy 0.9142231941223145\n",
            "Test set accuracy 0.917285144329071\n",
            "Epoch 5 in 1.02 sec\n",
            "Training set loss 0.5995732545852661\n",
            "Training set accuracy 0.9211103320121765\n",
            "Test set accuracy 0.9248046875\n",
            "Epoch 6 in 1.02 sec\n",
            "Training set loss 0.5955986976623535\n",
            "Training set accuracy 0.9261746406555176\n",
            "Test set accuracy 0.929394543170929\n",
            "Epoch 7 in 1.01 sec\n",
            "Training set loss 0.5924407839775085\n",
            "Training set accuracy 0.9303634762763977\n",
            "Test set accuracy 0.931835949420929\n",
            "Epoch 8 in 0.95 sec\n",
            "Training set loss 0.5899019241333008\n",
            "Training set accuracy 0.9341034889221191\n",
            "Test set accuracy 0.936328113079071\n",
            "Epoch 9 in 0.99 sec\n",
            "Training set loss 0.5878012180328369\n",
            "Training set accuracy 0.9369791746139526\n",
            "Test set accuracy 0.938769519329071\n",
            "Epoch 10 in 1.05 sec\n",
            "Training set loss 0.5860352516174316\n",
            "Training set accuracy 0.9397329092025757\n",
            "Test set accuracy 0.942187488079071\n",
            "Epoch 11 in 1.02 sec\n",
            "Training set loss 0.5845362544059753\n",
            "Training set accuracy 0.9419769048690796\n",
            "Test set accuracy 0.944140613079071\n",
            "Epoch 12 in 0.97 sec\n",
            "Training set loss 0.5832522511482239\n",
            "Training set accuracy 0.9445866346359253\n",
            "Test set accuracy 0.945996105670929\n",
            "Epoch 13 in 0.93 sec\n",
            "Training set loss 0.5821421146392822\n",
            "Training set accuracy 0.9467973709106445\n",
            "Test set accuracy 0.947460949420929\n",
            "Epoch 14 in 1.04 sec\n",
            "Training set loss 0.5811706781387329\n",
            "Training set accuracy 0.9483432769775391\n",
            "Test set accuracy 0.948437511920929\n",
            "Epoch 15 in 1.04 sec\n",
            "Training set loss 0.5803174376487732\n",
            "Training set accuracy 0.9501218199729919\n",
            "Test set accuracy 0.9498047232627869\n",
            "Epoch 16 in 0.99 sec\n",
            "Training set loss 0.5795542597770691\n",
            "Training set accuracy 0.9512022733688354\n",
            "Test set accuracy 0.951171875\n",
            "Epoch 17 in 1.06 sec\n",
            "Training set loss 0.5788754820823669\n",
            "Training set accuracy 0.952698290348053\n",
            "Test set accuracy 0.952832043170929\n",
            "Epoch 18 in 1.01 sec\n",
            "Training set loss 0.5782554149627686\n",
            "Training set accuracy 0.9542940258979797\n",
            "Test set accuracy 0.9541992545127869\n",
            "Epoch 19 in 1.00 sec\n",
            "Training set loss 0.5776922106742859\n",
            "Training set accuracy 0.9555240869522095\n",
            "Test set accuracy 0.95556640625\n"
          ]
        }
      ],
      "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",
        "    x = jnp.reshape(x, (len(x), NUM_PIXELS))\n",
        "    y = one_hot(y, NUM_LABELS)\n",
        "    x = jax.device_put(x, sharding)\n",
        "    y = jax.device_put(y, sharding)\n",
        "    params = jax.device_put(params, sharding.replicate())\n",
        "    params, loss_value = update(params, x, y, epoch)\n",
        "    losses.append(jnp.sum(loss_value))\n",
        "  epoch_time = time.time() - start_time\n",
        "\n",
        "  train_acc = accuracy(params, train_data)\n",
        "  test_acc = accuracy(params, test_data)\n",
        "  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
        "  print(\"Training set loss {}\".format(jnp.mean(jnp.array(losses))))\n",
        "  print(\"Training set accuracy {}\".format(train_acc))\n",
        "  print(\"Test set accuracy {}\".format(test_acc))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.debug.visualize_array_sharding(params[0][0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "FZxrQ5OvhV7u",
        "outputId": "5457df85-c04d-4412-d132-0a55ca2626cc"
      },
      "execution_count": 184,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────────────────────┐\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│        TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m        │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "└───────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────────────────┐\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│        TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>        │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "│                                   │\n",
              "└───────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### 4-way data parallelism, 2-way tensor parallelism"
      ],
      "metadata": {
        "id": "alhuQ3jN3Zp8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sharding = PositionalSharding(jax.devices()).reshape(4, 2)"
      ],
      "metadata": {
        "id": "vLmoPdbu3Y8p"
      },
      "execution_count": 185,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "LAYER_SIZES = [28*28, 10000, 10000, 10]\n",
        "PARAM_SCALE = 0.01\n",
        "\n",
        "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": "jyVcAxPC3mGG"
      },
      "execution_count": 186,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sharded_params = []\n",
        "for i,(w,b) in enumerate(init_params):\n",
        "  print(i, w.shape, b.shape)\n",
        "  if i==0:\n",
        "    w = jax.device_put(w, sharding.replicate(0))\n",
        "    b = jax.device_put(b, sharding.replicate(0))\n",
        "  elif i==1:\n",
        "    w = jax.device_put(w, sharding.replicate(0))\n",
        "    b = jax.device_put(b, sharding.replicate(0))\n",
        "  elif i==2:\n",
        "    w = jax.device_put(w, sharding.replicate())\n",
        "    b = jax.device_put(b, sharding.replicate())\n",
        "  sharded_params.append((w,b))\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-5dPxHms4a5E",
        "outputId": "00f41938-9fa8-4794-b679-7ad6aaf9699c"
      },
      "execution_count": 187,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0 (10000, 784) (10000,)\n",
            "1 (10000, 10000) (10000,)\n",
            "2 (10, 10000) (10,)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for (w,b) in init_params:\n",
        "  jax.debug.visualize_array_sharding(w)\n",
        "  jax.debug.visualize_array_sharding(b)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "1mmndGm66zkx",
        "outputId": "7480949b-83d7-45ec-e23a-d9e87b1c49ee"
      },
      "execution_count": 188,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┐\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│ TPU \u001b[1;36m0\u001b[0m │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "└───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┐\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "│       │\n",
              "└───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┐\n",
              "│ TPU \u001b[1;36m0\u001b[0m │\n",
              "└───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┐\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │\n",
              "└───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────────┐\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│         TPU \u001b[1;36m0\u001b[0m         │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "└───────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────┐\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│         TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>         │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "└───────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┐\n",
              "│ TPU \u001b[1;36m0\u001b[0m │\n",
              "└───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┐\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │\n",
              "└───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                    TPU \u001b[1;36m0\u001b[0m                                     │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                    TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>                                     │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────┐\n",
              "│ TPU \u001b[1;36m0\u001b[0m │\n",
              "└───────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────┐\n",
              "│ TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span> │\n",
              "└───────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for (w,b) in init_params:\n",
        "  jax.debug.visualize_array_sharding(jax.device_put(w, sharding.replicate()))\n",
        "  jax.debug.visualize_array_sharding(jax.device_put(b, sharding.replicate()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "vJaVTiEP8V43",
        "outputId": "1a4b282d-f803-4e7f-84f6-49128c5c2b06"
      },
      "execution_count": 189,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────┐\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "└───────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────┐\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "│                   │\n",
              "└───────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────────┐\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│  TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m  │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "└───────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────────┐\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│  TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>  │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "│                       │\n",
              "└───────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                             TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                             TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for (w,b) in sharded_params:\n",
        "  jax.debug.visualize_array_sharding(w)\n",
        "  jax.debug.visualize_array_sharding(b)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "o28ff-t966yS",
        "outputId": "dba7c757-2054-4915-e904-03f93d134ce2"
      },
      "execution_count": 190,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                             TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌──────────────────────────────────────────────────────────────────────────────┐\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                             TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "│                                                                              │\n",
              "└──────────────────────────────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────────────┐\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m1\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m6\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "└───────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────────────┐\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "└───────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "params = sharded_params\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",
        "    x = jax.device_put(x, sharding.replicate(1))\n",
        "    y = jax.device_put(y, sharding.replicate(1))\n",
        "    #params = jax.device_put(params, sharding.replicate())\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))\n",
        "\n",
        "  jax.debug.visualize_array_sharding(params[0][0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "9y4hEhkp7GQI",
        "outputId": "229e456d-c65e-4f72-c9db-fd3ddaf41de6"
      },
      "execution_count": 191,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0 in 6.61 sec\n",
            "Training set loss 0.6637932062149048\n",
            "Training set accuracy 0.8757756948471069\n",
            "Test set accuracy 0.8863281607627869\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1 in 3.96 sec\n",
            "Training set loss 0.6224857568740845\n",
            "Training set accuracy 0.8902149200439453\n",
            "Test set accuracy 0.8956055045127869\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 2 in 3.95 sec\n",
            "Training set loss 0.6166571378707886\n",
            "Training set accuracy 0.8960660099983215\n",
            "Test set accuracy 0.900097668170929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 3 in 3.96 sec\n",
            "Training set loss 0.613541841506958\n",
            "Training set accuracy 0.8999667167663574\n",
            "Test set accuracy 0.9034180045127869\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 4 in 3.97 sec\n",
            "Training set loss 0.6112023591995239\n",
            "Training set accuracy 0.9033576846122742\n",
            "Test set accuracy 0.907421886920929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 5 in 3.95 sec\n",
            "Training set loss 0.6091800332069397\n",
            "Training set accuracy 0.9063829779624939\n",
            "Test set accuracy 0.909375011920929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 6 in 3.95 sec\n",
            "Training set loss 0.6072940230369568\n",
            "Training set accuracy 0.9090425372123718\n",
            "Test set accuracy 0.911816418170929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 7 in 3.95 sec\n",
            "Training set loss 0.6054378747940063\n",
            "Training set accuracy 0.911436140537262\n",
            "Test set accuracy 0.9150390625\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 8 in 3.95 sec\n",
            "Training set loss 0.6035282611846924\n",
            "Training set accuracy 0.9151983261108398\n",
            "Test set accuracy 0.918749988079071\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 9 in 3.94 sec\n",
            "Training set loss 0.6015031337738037\n",
            "Training set accuracy 0.9181182980537415\n",
            "Test set accuracy 0.921191394329071\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 10 in 3.94 sec\n",
            "Training set loss 0.5993688702583313\n",
            "Training set accuracy 0.9215425252914429\n",
            "Test set accuracy 0.9237304925918579\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 11 in 3.93 sec\n",
            "Training set loss 0.5972492694854736\n",
            "Training set accuracy 0.9245678186416626\n",
            "Test set accuracy 0.9267578125\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 12 in 3.94 sec\n",
            "Training set loss 0.5953165888786316\n",
            "Training set accuracy 0.9276872873306274\n",
            "Test set accuracy 0.9283203482627869\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 13 in 3.96 sec\n",
            "Training set loss 0.5936313271522522\n",
            "Training set accuracy 0.9301806092262268\n",
            "Test set accuracy 0.930371105670929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 14 in 3.95 sec\n",
            "Training set loss 0.5921530723571777\n",
            "Training set accuracy 0.932541012763977\n",
            "Test set accuracy 0.932812511920929\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 15 in 3.96 sec\n",
            "Training set loss 0.590818464756012\n",
            "Training set accuracy 0.9347129464149475\n",
            "Test set accuracy 0.934374988079071\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 16 in 3.96 sec\n",
            "Training set loss 0.5895960330963135\n",
            "Training set accuracy 0.9365746378898621\n",
            "Test set accuracy 0.936328113079071\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 17 in 3.93 sec\n",
            "Training set loss 0.5884610414505005\n",
            "Training set accuracy 0.9382202625274658\n",
            "Test set accuracy 0.9383789300918579\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 18 in 3.93 sec\n",
            "Training set loss 0.5874113440513611\n",
            "Training set accuracy 0.9402315616607666\n",
            "Test set accuracy 0.9404296875\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 19 in 3.93 sec\n",
            "Training set loss 0.5864334106445312\n",
            "Training set accuracy 0.9416278004646301\n",
            "Test set accuracy 0.941699206829071\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU \u001b[1;36m0\u001b[0m,\u001b[1;36m2\u001b[0m,\u001b[1;36m4\u001b[0m,\u001b[1;36m6\u001b[0m│TPU \u001b[1;36m1\u001b[0m,\u001b[1;36m3\u001b[0m,\u001b[1;36m5\u001b[0m,\u001b[1;36m7\u001b[0m│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┌───────────┬───────────┐\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6</span>│TPU <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5</span>,<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>│\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "│           │           │\n",
              "└───────────┴───────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "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": [],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}