Chapter-8/JAX_in_Action_Chapter_8_Tensor_Sharding.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pXcZUT7WJv4e"
      },
      "source": [
        "## Using Cloud TPU and a Local runtime in Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ps5o9t18KCb3"
      },
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DECNlXiNKuc2",
        "outputId": "e8cf6e70-1357-496d-aa30-95506609e5b4"
      },
      "outputs": [],
      "source": [
        "!pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-GXV0DH2P7kf"
      },
      "source": [
        "Needed for pretty sharding visualization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ryHjnoA7P5pQ",
        "outputId": "6e9b1fc6-cefb-4d71-8447-b81ac60340f5"
      },
      "outputs": [],
      "source": [
        "!pip install rich"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "9z7CUfFvJ7JN"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_RJ5XZCcKAYD",
        "outputId": "06470233-f45e-4085-8b53-7dafdfd10bfd"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2hICp5VhC6Bt",
        "outputId": "ecb48f78-7eb3-4c36-b36e-9498defbee2a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'0.4.31'"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XRLNPE2Ese1D"
      },
      "source": [
        "## Using Tensor Sharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "e3HsZ2PIsl3t"
      },
      "outputs": [],
      "source": [
        "from jax.experimental import mesh_utils\n",
        "from jax.sharding import PositionalSharding, NamedSharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "-Bocy386sr3n"
      },
      "outputs": [],
      "source": [
        "from jax import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "SI0alZjP56SZ"
      },
      "outputs": [],
      "source": [
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1wKd0cB8iu5"
      },
      "source": [
        "### Dot example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "2p2R6uG5tVpV"
      },
      "outputs": [],
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "SDMWf5eIswdB"
      },
      "outputs": [],
      "source": [
        "rng_key = random.PRNGKey(42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "THq8XQoAs8CH",
        "outputId": "4a2d4210-5c82-4aa0-a4f2-bee282d41628"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((4000, 10000), (4000, 10000))"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "vs = random.normal(rng_key, shape=(8_000,10_000))\n",
        "v1s = vs[:4_000,:]\n",
        "v2s = vs[4_000:,:]\n",
        "\n",
        "v1s.shape, v2s.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "ap2c7J2lPz-H",
        "outputId": "7064d552-647e-41e4-b39f-87b946dab584"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v1s)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "we0gE4dV9bHv"
      },
      "source": [
        "### Positional sharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "D6OxkmCbQeJw"
      },
      "outputs": [],
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((4,1)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C86DIH_uQ7DT",
        "outputId": "1b2a7d3e-08eb-41f9-90d0-c220296161d6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 0}]\n",
              "                    [{TPU 2}]\n",
              "                    [{TPU 1}]\n",
              "                    [{TPU 3}]], shape=(4, 1))"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "6YuoF1hhRgRs"
      },
      "outputs": [],
      "source": [
        "v1sp = jax.device_put(v1s, sharding)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "jaxlib.xla_extension.ArrayImpl"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(v1s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3ASNkCAddFkB",
        "outputId": "3a1ac310-78ac-4691-9b07-8d99271e7d65"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "jaxlib.xla_extension.ArrayImpl"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "DRm5TQDrYwqo",
        "outputId": "e6a3b99a-d696-47a4-e571-438028703898"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                            TPU 2                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                            TPU 1                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                            TPU 3                             </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                            \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                             \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "9N2J_H_KRRc6"
      },
      "outputs": [],
      "source": [
        "v2sp = jax.device_put(v2s, sharding)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "0_5ZNoMZQhrZ",
        "outputId": "b4cbcf66-f46b-43ae-b190-36e4edbb5f36"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                            TPU 2                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                            TPU 1                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                            TPU 3                             </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                            \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                             \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rAe6pb2wRhha"
      },
      "source": [
        "Input is sharded across all the devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "qjjU9AGsu9cX"
      },
      "outputs": [],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vnl0vIl3vAOc",
        "outputId": "9dadfa81-0d10-4561-c44b-4686a4317931"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "d.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "lM6bu5xqvFlu",
        "outputId": "2ef18cfd-c9a7-47e1-9f93-f9e64a2ae16c"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 2  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 1  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r6LQ3H0xvV-S",
        "outputId": "e57a9bf2-bba4-439e-dc56-4602d8a88942"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "570 μs ± 3.61 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jax.vmap(dot)(v1sp, v2sp).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rdbMpGTkvcrw",
        "outputId": "eaf3f7df-2a59-4aba-ccdb-b2462fa0c2ad"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "738 μs ± 2.69 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jax.vmap(dot)(v1s, v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ovKaom8nvoN-",
        "outputId": "8d6e80b7-3709-439d-cfb1-33f54407df95"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4000,10000] b:f32[4000,10000]. let\n",
              "    c:f32[4000] = dot_general[\n",
              "      dimension_numbers=(([1], [1]), ([0], [0]))\n",
              "      preferred_element_type=float32\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 24,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(jax.vmap(dot))(v1s, v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9MLQkQxBv0LP",
        "outputId": "66b16e90-5dac-46ad-e55d-c1f5c1e5298c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "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_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(128)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(128)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(128)} add(f32[]{:T(128)} %scalar_lhs, f32[]{:T(128)} %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(128)} constant(0)\n",
            "  ROOT %reduce.1 = f32[4000]{0:T(1024)} reduce(f32[4000,10000]{1,0:T(8,128)} %multiply.1, f32[]{:T(128)} %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=float32]\" source_file=\"/tmp/ipykernel_972295/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), metadata={op_name=\"v2\"}\n",
            "  %Arg_0.1 = f32[4000,10000]{1,0:T(8,128)} parameter(0), metadata={op_name=\"v1\"}\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=float32]\" source_file=\"/tmp/ipykernel_972295/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"500\",\"1\"],\"input_window_bounds\":[],\"estimated_cycles\":\"376370\",\"iteration_bounds\":[\"1\",\"79\"]},\"megacore_config\":{\"megacore_split_dim\":\"1\",\"megacore_allreduce_bytes\":\"16384\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"8208384\"}],\"retry_config\":{\"retry_count\":\"0\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1s, v2s).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HfVoi5Z72lUG",
        "outputId": "61032fcf-266a-4e1c-8e05-40ea521549e1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{ lambda ; a:f32[4000,10000] b:f32[4000,10000]. let\n",
              "    c:f32[4000] = dot_general[\n",
              "      dimension_numbers=(([1], [1]), ([0], [0]))\n",
              "      preferred_element_type=float32\n",
              "    ] a b\n",
              "  in (c,) }"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.make_jaxpr(jax.vmap(dot))(v1sp, v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uRElWOlm2o57",
        "outputId": "564ada03-cd38-40e0-816c-0fc278bb84f1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[1000,10000]{1,0:T(8,128)}, f32[1000,10000]{1,0:T(8,128)})->f32[1000]{0:T(1024)}}, allow_spmd_sharding_propagation_to_parameters={false,false}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=4\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(128)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(128)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(128)} add(f32[]{:T(128)} %scalar_lhs, f32[]{:T(128)} %scalar_rhs)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.2: f32[1000,10000], param_1.2: f32[1000,10000]) -> f32[1000] {\n",
            "  %param_0.2 = f32[1000,10000]{1,0:T(8,128)} parameter(0)\n",
            "  %param_1.2 = f32[1000,10000]{1,0:T(8,128)} parameter(1)\n",
            "  %multiply.2 = f32[1000,10000]{1,0:T(8,128)} multiply(f32[1000,10000]{1,0:T(8,128)} %param_0.2, f32[1000,10000]{1,0:T(8,128)} %param_1.2)\n",
            "  %constant.2 = f32[]{:T(128)} constant(0)\n",
            "  ROOT %reduce.2 = f32[1000]{0:T(1024)} reduce(f32[1000,10000]{1,0:T(8,128)} %multiply.2, f32[]{:T(128)} %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=float32]\" source_file=\"/tmp/ipykernel_972295/2294503736.py\" source_line=2}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[1000,10000], param.1: f32[1000,10000]) -> f32[1000] {\n",
            "  %param.1 = f32[1000,10000]{1,0:T(8,128)} parameter(1), sharding={devices=[4,1]0,1,2,3}, metadata={op_name=\"v2\"}\n",
            "  %param = f32[1000,10000]{1,0:T(8,128)} parameter(0), sharding={devices=[4,1]0,1,2,3}, metadata={op_name=\"v1\"}\n",
            "  ROOT %fusion = f32[1000]{0:T(1024)} fusion(f32[1000,10000]{1,0:T(8,128)} %param, f32[1000,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=float32]\" source_file=\"/tmp/ipykernel_972295/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"125\",\"5\"],\"input_window_bounds\":[],\"estimated_cycles\":\"94648\",\"iteration_bounds\":[\"1\",\"16\"]},\"megacore_config\":{\"megacore_split_dim\":\"1\",\"megacore_allreduce_bytes\":\"4096\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"10244096\"}],\"retry_config\":{\"retry_count\":\"0\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1sp, v2sp).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_y63HyZIeswz"
      },
      "source": [
        "### 2D mesh example\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "UiCNiPlixOW8"
      },
      "outputs": [],
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,2)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "_Ts2Yx3rxWui"
      },
      "outputs": [],
      "source": [
        "v1sp = jax.device_put(v1s, sharding)\n",
        "v2sp = jax.device_put(v2s, sharding)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "PT7eapRUxiXZ",
        "outputId": "08595b8c-c364-435c-dc22-37ef67130780"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">             TPU 0             </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">             TPU 1             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">             TPU 2             </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">             TPU 3             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m             \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m             \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m             \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m             \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u8ZPzKbfyGE6",
        "outputId": "5def7ebf-2bde-4da3-d1f1-b51ca5d96cf2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)\n",
        "d.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "pflrYiYLyMPQ",
        "outputId": "facdaf00-05a9-4653-fde7-6bae8bbb0a1b"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> TPU 0,1  TPU 2,3 </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fiuoqoK2yRhB",
        "outputId": "eccdbb82-91c4-45e5-a7ac-85a922e2f717"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "583 μs ± 6.7 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jax.vmap(dot)(v1sp, v2sp).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c_zOa1uyyR3r",
        "outputId": "76f2cf0c-2c63-4275-d3c6-7cf51cb906de"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "741 μs ± 5.48 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jax.vmap(dot)(v1s, v2s).block_until_ready()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1XOdpG2MA0Wv"
      },
      "source": [
        "Looking at HLO (note the all-reduce operation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fGkNyNn52s8T",
        "outputId": "49ac0346-fd7e-41b1-b188-019a66146523"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[2000,5000]{1,0:T(8,128)}, f32[2000,5000]{1,0:T(8,128)})->f32[2000]{0:T(1024)}}, allow_spmd_sharding_propagation_to_parameters={false,false}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=4\n",
            "\n",
            "%scalar_add_computation.clone (scalar_lhs.1: f32[], scalar_rhs.1: f32[]) -> f32[] {\n",
            "  %scalar_rhs.1 = f32[]{:T(128)} parameter(1)\n",
            "  %scalar_lhs.1 = f32[]{:T(128)} parameter(0)\n",
            "  ROOT %add.1 = f32[]{:T(128)} add(f32[]{:T(128)} %scalar_lhs.1, f32[]{:T(128)} %scalar_rhs.1)\n",
            "}\n",
            "\n",
            "%scalar_add_computation (scalar_lhs: f32[], scalar_rhs: f32[]) -> f32[] {\n",
            "  %scalar_rhs = f32[]{:T(128)} parameter(1)\n",
            "  %scalar_lhs = f32[]{:T(128)} parameter(0)\n",
            "  ROOT %add = f32[]{:T(128)} add(f32[]{:T(128)} %scalar_lhs, f32[]{:T(128)} %scalar_rhs)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0.2: f32[2000,5000], param_1.2: f32[2000,5000]) -> f32[2000] {\n",
            "  %param_0.2 = f32[2000,5000]{1,0:T(8,128)} parameter(0)\n",
            "  %param_1.2 = f32[2000,5000]{1,0:T(8,128)} parameter(1)\n",
            "  %multiply.2 = f32[2000,5000]{1,0:T(8,128)} multiply(f32[2000,5000]{1,0:T(8,128)} %param_0.2, f32[2000,5000]{1,0:T(8,128)} %param_1.2)\n",
            "  %constant.2 = f32[]{:T(128)} constant(0)\n",
            "  ROOT %reduce.2 = f32[2000]{0:T(1024)S(3)} reduce(f32[2000,5000]{1,0:T(8,128)} %multiply.2, f32[]{:T(128)} %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=float32]\" source_file=\"/tmp/ipykernel_972295/2294503736.py\" source_line=2}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[2000,5000], param.1: f32[2000,5000]) -> f32[2000] {\n",
            "  %param.1 = f32[2000,5000]{1,0:T(8,128)} parameter(1), sharding={devices=[2,2]0,1,2,3}, metadata={op_name=\"v2\"}\n",
            "  %param = f32[2000,5000]{1,0:T(8,128)} parameter(0), sharding={devices=[2,2]0,1,2,3}, metadata={op_name=\"v1\"}\n",
            "  %fusion = f32[2000]{0:T(1024)S(3)} fusion(f32[2000,5000]{1,0:T(8,128)} %param, f32[2000,5000]{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=float32]\" source_file=\"/tmp/ipykernel_972295/2294503736.py\" source_line=2}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[],\"output_window_bounds\":[\"250\",\"2\"],\"input_window_bounds\":[],\"estimated_cycles\":\"95280\",\"iteration_bounds\":[\"1\",\"20\"]},\"megacore_config\":{\"megacore_split_dim\":\"1\",\"megacore_allreduce_bytes\":\"8192\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"8200192\"}],\"retry_config\":{\"retry_count\":\"0\"}}\n",
            "  ROOT %all-reduce = f32[2000]{0:T(1024)} all-reduce(f32[2000]{0:T(1024)S(3)} %fusion), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, to_apply=%scalar_add_computation.clone, metadata={op_name=\"jit(dot)/jit(main)/dot_general[dimension_numbers=(((1,), (1,)), ((0,), (0,))) precision=None preferred_element_type=float32]\" source_file=\"/tmp/ipykernel_972295/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\"},\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"49152\"}],\"retry_config\":{\"retry_count\":\"0\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = jax.jit(jax.vmap(dot)).lower(v1sp, v2sp).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "516tLncUA2Z5"
      },
      "source": [
        "### Using replication"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 128,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "ooOHrGv64TZF",
        "outputId": "539aa17a-df53-4186-9e00-39c5a62d53d0"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                           TPU 0,1                            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                           TPU 2,3                            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                           \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                           \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,2)))\n",
        "\n",
        "v1sp = jax.device_put(v1s, sharding.replicate(axis=1))\n",
        "\n",
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 129,
      "metadata": {
        "id": "Y1wPyLI1C6P3"
      },
      "outputs": [],
      "source": [
        "A = random.normal(rng_key, shape=(10000,2000))\n",
        "B = random.normal(rng_key, shape=(2000,5000))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 134,
      "metadata": {
        "id": "qFct_kR6DFho"
      },
      "outputs": [],
      "source": [
        "Ad = jax.device_put(A, sharding.replicate(0))\n",
        "Bd = jax.device_put(B, sharding.replicate())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 135,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "OYhgu1U4DVYg",
        "outputId": "68ffb6bb-d1f1-4d27-be31-f2a58090ab61"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\"> TPU 0,2  TPU 1,3 </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                  </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,2\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 1,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         TPU 0,1,2,3                          </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                          \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(Ad)\n",
        "jax.debug.visualize_array_sharding(Bd)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 136,
      "metadata": {
        "id": "ld4UAh1QDOxU"
      },
      "outputs": [],
      "source": [
        "Cd = jnp.dot(Ad, Bd)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 137,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "QdzAVah9DXyA",
        "outputId": "0b002581-f0dc-4da9-bb46-3a95d3bb2853"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3 </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(Cd)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 45,
      "metadata": {
        "id": "_W-UMoSkDZA2"
      },
      "outputs": [],
      "source": [
        "C = A @ B"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MNHr5cE8DcS2",
        "outputId": "4c7571a8-8dcb-4b98-f302-bf756df8e40f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(False, dtype=bool)"
            ]
          },
          "execution_count": 46,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.numpy.array_equal(C,Cd)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "RI9dX5LlE7sh",
        "outputId": "d1878349-77f6-449c-9ff2-36ce6c3c98f9"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   TPU 0    </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(C)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DF3Gv28IVO6V",
        "outputId": "7c51e884-2ba2-46af-ec9e-e684e97c93d5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "((10000, 5000), (10000, 5000))"
            ]
          },
          "execution_count": 48,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "C.shape, Cd.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RtVhE7KKVeji",
        "outputId": "e1222f9f-09e7-48f6-a036-65d22cfb62e2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(Array(43.027626, dtype=float32), Array(43.027634, dtype=float32))"
            ]
          },
          "execution_count": 49,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "C[12,3], Cd[12,3]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "QxgSv_ZSVquX"
      },
      "outputs": [],
      "source": [
        "Ca = jnp.dot(A, B)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xO_sxUFrVtpN",
        "outputId": "91b814fd-dd30-43cf-c98d-e603ffa3feb9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "execution_count": 51,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.numpy.array_equal(C,Ca)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "wQvIChAAV0f_",
        "outputId": "5ec8e7f5-04d3-444c-ba7e-893abb4fef4e"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   TPU 0    </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(Ca)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "id": "UdXSxOhHV7wl"
      },
      "outputs": [],
      "source": [
        "d = (Cd - C)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I-uXR9-EV_A1",
        "outputId": "5f3937e4-cd16-4b27-cfcf-d5232a441443"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(Array(6.1035156e-05, dtype=float32), Array(0.01788047, dtype=float32))"
            ]
          },
          "execution_count": 54,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.max(d), jnp.sum(d)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6RCR-BcuWtvi",
        "outputId": "31f53d02-e4ce-4db3-9038-e994292b1634"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "400 μs ± 2.26 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jnp.dot(Ad, Bd).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GSS3YK4IW21p",
        "outputId": "0510c115-6b9f-4673-8550-de7d59acb2c7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "958 μs ± 4.32 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit jnp.dot(A, B).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fF_EaKAGW5Nk",
        "outputId": "290105b2-8a8d-4a16-ab0e-7f4965b91afe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "956 μs ± 3.77 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit (A@B).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 58,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Nrq6pY5fH3Id",
        "outputId": "af8db4e2-6f36-4184-a9e6-d6885c4a76dd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "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_parameters={true,true}, 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)S(3)} parameter(0)\n",
            "  ROOT %bitcast = f32[10000,2000]{0,1:T(8,128)} bitcast(f32[10000,2000]{0,1:T(8,128)S(3)} %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)S(3)} parameter(0)\n",
            "  %fusion.1 = f32[10000,2000]{0,1:T(8,128)} fusion(f32[10000,2000]{0,1:T(8,128)S(3)} %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=float32]\" source_file=\"/tmp/ipykernel_972295/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_0.1 = f32[10000,2000]{0,1:T(8,128)} parameter(0), metadata={op_name=\"a\"}\n",
            "  %copy-start = (f32[10000,2000]{0,1:T(8,128)S(3)}, f32[10000,2000]{0,1:T(8,128)}, u32[]{:S(2)}) copy-start(f32[10000,2000]{0,1:T(8,128)} %Arg_0.1), cross_program_prefetch_index=0\n",
            "  %Arg_1.2 = f32[2000,5000]{1,0:T(8,128)} parameter(1), metadata={op_name=\"b\"}\n",
            "  %copy-done = f32[10000,2000]{0,1:T(8,128)S(3)} copy-done((f32[10000,2000]{0,1:T(8,128)S(3)}, f32[10000,2000]{0,1:T(8,128)}, u32[]{:S(2)}) %copy-start)\n",
            "  ROOT %fusion = f32[10000,5000]{0,1:T(8,128)} fusion(f32[10000,2000]{0,1:T(8,128)S(3)} %copy-done, 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=float32]\" source_file=\"/tmp/ipykernel_972295/1507742915.py\" source_line=1}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"125\",\"5\"],\"output_window_bounds\":[\"80\",\"8\"],\"input_window_bounds\":[\"125\",\"8\"],\"estimated_cycles\":\"943119\",\"iteration_bounds\":[\"8\",\"10\",\"2\"]},\"megacore_config\":{\"megacore_split_dim\":\"0\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"12427264\"}],\"retry_config\":{\"retry_count\":\"0\"},\"convolution_algorithm_config\":{\"emitter\":\"EmitInputBatchInLanes\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = jnp.dot.lower(A, B).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 59,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P1auVbzQHxAG",
        "outputId": "54477235-7c3a-497b-c9db-f67293e20a14"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HloModule jit_dot, is_scheduled=true, entry_computation_layout={(f32[5000,2000]{1,0:T(8,128)}, f32[2000,2500]{1,0:T(8,128)})->f32[5000,2500]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={false,false}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=4\n",
            "\n",
            "%bitcast_fusion (bf16input: f32[5000,2000]) -> f32[5000,2000] {\n",
            "  %bf16input = f32[5000,2000]{1,0:T(8,128)S(3)} parameter(0)\n",
            "  ROOT %bitcast = f32[5000,2000]{1,0:T(8,128)} bitcast(f32[5000,2000]{1,0:T(8,128)S(3)} %bf16input)\n",
            "}\n",
            "\n",
            "%bitcast_fusion.1 (bf16input.1: f32[2000,2500]) -> f32[2000,2500] {\n",
            "  %bf16input.1 = f32[2000,2500]{1,0:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast.1 = f32[2000,2500]{1,0:T(8,128)} bitcast(f32[2000,2500]{1,0:T(8,128)} %bf16input.1)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0: f32[5000,2000], param_1: f32[2000,2500]) -> f32[5000,2500] {\n",
            "  %param_0 = f32[5000,2000]{1,0:T(8,128)S(3)} parameter(0)\n",
            "  %fusion.1 = f32[5000,2000]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)S(3)} %param_0), kind=kLoop, calls=%bitcast_fusion\n",
            "  %param_1 = f32[2000,2500]{1,0:T(8,128)} parameter(1)\n",
            "  %fusion.2 = f32[2000,2500]{1,0:T(8,128)} fusion(f32[2000,2500]{1,0:T(8,128)} %param_1), kind=kLoop, calls=%bitcast_fusion.1\n",
            "  ROOT %convolution.1 = f32[5000,2500]{1,0:T(8,128)} convolution(f32[5000,2000]{1,0:T(8,128)} %fusion.1, f32[2000,2500]{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=float32]\" source_file=\"/tmp/ipykernel_972295/1507742915.py\" source_line=1}\n",
            "}\n",
            "\n",
            "ENTRY %main.4_spmd (param: f32[5000,2000], param.1: f32[2000,2500]) -> f32[5000,2500] {\n",
            "  %param = f32[5000,2000]{1,0:T(8,128)} parameter(0), sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}, metadata={op_name=\"a\"}\n",
            "  %copy-start = (f32[5000,2000]{1,0:T(8,128)S(3)}, f32[5000,2000]{1,0:T(8,128)}, u32[]{:S(2)}) copy-start(f32[5000,2000]{1,0:T(8,128)} %param), cross_program_prefetch_index=0\n",
            "  %param.1 = f32[2000,2500]{1,0:T(8,128)} parameter(1), sharding={devices=[1,2,2]0,2,1,3 last_tile_dim_replicate}, metadata={op_name=\"b\"}\n",
            "  %copy-done = f32[5000,2000]{1,0:T(8,128)S(3)} copy-done((f32[5000,2000]{1,0:T(8,128)S(3)}, f32[5000,2000]{1,0:T(8,128)}, u32[]{:S(2)}) %copy-start)\n",
            "  ROOT %fusion = f32[5000,2500]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)S(3)} %copy-done, f32[2000,2500]{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=float32]\" source_file=\"/tmp/ipykernel_972295/1507742915.py\" source_line=1}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"250\",\"2\"],\"output_window_bounds\":[\"209\",\"2\"],\"input_window_bounds\":[\"209\",\"16\"],\"estimated_cycles\":\"238762\",\"iteration_bounds\":[\"10\",\"3\",\"1\"]},\"megacore_config\":{\"megacore_split_dim\":\"0\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"7520256\"}],\"retry_config\":{\"retry_count\":\"0\"},\"convolution_algorithm_config\":{\"emitter\":\"EmitAllBatchInSublanes\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = jnp.dot.lower(Ad, Bd).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "USf1avNxJ8Na"
      },
      "source": [
        "### Using sharding constraints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "id": "GquL4IMFJ-rd"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "id": "7LEVlx4akw5c"
      },
      "outputs": [],
      "source": [
        "sharding = PositionalSharding(mesh_utils.create_device_mesh((2,2)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 62,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "G1ThCrDikrjL",
        "outputId": "70966bbc-d5e2-45ba-9e51-9f889a591850"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(A)\n",
        "jax.debug.visualize_array_sharding(B)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "metadata": {
        "id": "HiJOk9HEk7kT"
      },
      "outputs": [],
      "source": [
        "d = distributed_mul(A, B, sharding)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 64,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "knUitJvSk_1Z",
        "outputId": "37c86573-0a41-458e-d94f-60bab8e578c8"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fkhWKJZBnPC0",
        "outputId": "75b0090d-d5bc-4bff-aec9-945705258aba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10000, 5000)"
            ]
          },
          "execution_count": 65,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "d.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "metadata": {
        "id": "_ArK0PqwlgOj"
      },
      "outputs": [],
      "source": [
        "@jit\n",
        "def nondistributed_mul(a, b):\n",
        "  return jnp.dot(a, b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {
        "id": "EGSWUis-liuu"
      },
      "outputs": [],
      "source": [
        "dn = nondistributed_mul(A, B)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "PrAZODtAllOl",
        "outputId": "449d3c52-b782-47ee-91e4-caf6d391ad6e"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   TPU 0    </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(dn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JIo3eJqklrmG",
        "outputId": "f1bc06a7-e3d3-48a9-cdeb-d665340f8318"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "3.99 ms ± 33.1 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit distributed_mul(A, B, sharding).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wmyvZXW2luce",
        "outputId": "bfb69e03-8251-4dd1-951d-301c270ca23e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "953 μs ± 351 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
          ]
        }
      ],
      "source": [
        "%timeit nondistributed_mul(A, B).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YXUu0WcaHf0H",
        "outputId": "f2b131d0-b416-4450-a61a-b5e7c9b05b95"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "HloModule jit_distributed_mul, is_scheduled=true, entry_computation_layout={(f32[5000,2000]{1,0:T(8,128)}, f32[2000,2500]{1,0:T(8,128)})->f32[5000,2500]{1,0:T(8,128)}}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=4\n",
            "\n",
            "%bitcast_fusion (bf16input: f32[5000,2000]) -> f32[5000,2000] {\n",
            "  %bf16input = f32[5000,2000]{1,0:T(8,128)S(3)} parameter(0)\n",
            "  ROOT %bitcast = f32[5000,2000]{1,0:T(8,128)} bitcast(f32[5000,2000]{1,0:T(8,128)S(3)} %bf16input)\n",
            "}\n",
            "\n",
            "%bitcast_fusion.1 (bf16input.1: f32[2000,2500]) -> f32[2000,2500] {\n",
            "  %bf16input.1 = f32[2000,2500]{1,0:T(8,128)} parameter(0)\n",
            "  ROOT %bitcast.1 = f32[2000,2500]{1,0:T(8,128)} bitcast(f32[2000,2500]{1,0:T(8,128)} %bf16input.1)\n",
            "}\n",
            "\n",
            "%fused_computation (param_0: f32[5000,2000], param_1: f32[2000,2500]) -> f32[5000,2500] {\n",
            "  %param_0 = f32[5000,2000]{1,0:T(8,128)S(3)} parameter(0)\n",
            "  %fusion.1 = f32[5000,2000]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)S(3)} %param_0), kind=kLoop, calls=%bitcast_fusion\n",
            "  %param_1 = f32[2000,2500]{1,0:T(8,128)} parameter(1)\n",
            "  %fusion.2 = f32[2000,2500]{1,0:T(8,128)} fusion(f32[2000,2500]{1,0:T(8,128)} %param_1), kind=kLoop, calls=%bitcast_fusion.1\n",
            "  ROOT %convolution.1 = f32[5000,2500]{1,0:T(8,128)} convolution(f32[5000,2000]{1,0:T(8,128)} %fusion.1, f32[2000,2500]{1,0:T(8,128)} %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=float32]\" source_file=\"/tmp/ipykernel_972295/1742017423.py\" source_line=8}\n",
            "}\n",
            "\n",
            "ENTRY %main.6_spmd (param: f32[5000,2000], param.1: f32[2000,2500]) -> f32[5000,2500] {\n",
            "  %param = f32[5000,2000]{1,0:T(8,128)} parameter(0), sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}, metadata={op_name=\"a\"}\n",
            "  %copy-start = (f32[5000,2000]{1,0:T(8,128)S(3)}, f32[5000,2000]{1,0:T(8,128)}, u32[]{:S(2)}) copy-start(f32[5000,2000]{1,0:T(8,128)} %param), cross_program_prefetch_index=0\n",
            "  %param.1 = f32[2000,2500]{1,0:T(8,128)} parameter(1), sharding={devices=[1,2,2]0,2,1,3 last_tile_dim_replicate}, metadata={op_name=\"b\"}\n",
            "  %copy-done = f32[5000,2000]{1,0:T(8,128)S(3)} copy-done((f32[5000,2000]{1,0:T(8,128)S(3)}, f32[5000,2000]{1,0:T(8,128)}, u32[]{:S(2)}) %copy-start)\n",
            "  ROOT %fusion = f32[5000,2500]{1,0:T(8,128)} fusion(f32[5000,2000]{1,0:T(8,128)S(3)} %copy-done, f32[2000,2500]{1,0:T(8,128)} %param.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=float32]\" source_file=\"/tmp/ipykernel_972295/1742017423.py\" source_line=8}, backend_config={\"flag_configs\":[],\"window_config\":{\"kernel_window_bounds\":[\"250\",\"2\"],\"output_window_bounds\":[\"209\",\"2\"],\"input_window_bounds\":[\"209\",\"16\"],\"estimated_cycles\":\"238762\",\"iteration_bounds\":[\"10\",\"3\",\"1\"]},\"megacore_config\":{\"megacore_split_dim\":\"0\"},\"scoped_memory_configs\":[],\"used_scoped_memory_configs\":[{\"memory_space\":\"1\",\"offset\":\"0\",\"size\":\"7520256\"}],\"retry_config\":{\"retry_count\":\"0\"},\"convolution_algorithm_config\":{\"emitter\":\"EmitAllBatchInSublanes\"}}\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "modules = distributed_mul.lower(A, B, sharding).compile().as_text()\n",
        "print(modules)\n",
        "#for hlo in modules:\n",
        "#  print(hlo.to_string())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rPIpRTgO9fDa"
      },
      "source": [
        "### Named sharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "id": "BbfZ4Zto9gLY"
      },
      "outputs": [],
      "source": [
        "from jax.sharding import Mesh\n",
        "from jax.sharding import PartitionSpec as P\n",
        "from jax.sharding import NamedSharding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "id": "Pz42EU019j99"
      },
      "outputs": [],
      "source": [
        "mesh = Mesh(mesh_utils.create_device_mesh((2,2)), axis_names=('batch', 'features'))\n",
        "sharding = NamedSharding(mesh, P('batch', 'features'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "iu1O52k5-caD",
        "outputId": "4ad2e8b0-d672-4489-c6c6-14d9a3557d5c"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">             TPU 0             </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">             TPU 1             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                               </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">             TPU 2             </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">             TPU 3             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                               </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                               </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m             \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                               \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m             \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m             \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m             \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                               \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                               \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "v1sp = jax.device_put(v1s, sharding)\n",
        "v2sp = jax.device_put(v2s, sharding)\n",
        "\n",
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "18uhGU7pAgui",
        "outputId": "aed16b54-a4ed-4aa2-f057-30ce3ea74a78"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(4000,)"
            ]
          },
          "execution_count": 77,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)\n",
        "d.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RBjyztjhKLxx"
      },
      "source": [
        "### Device placement policy and errors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zCpzsk5vbNne"
      },
      "source": [
        "Different devices:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "id": "qqIQJ2ejKJw9"
      },
      "outputs": [],
      "source": [
        "sharding_a = PositionalSharding(np.array(jax.devices()[:2]).reshape(2,1))\n",
        "sharding_b = PositionalSharding(np.array(jax.devices()[2:]).reshape(2,1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JAbiGQEUaE4d",
        "outputId": "45fa969a-47a3-401f-e355-59ba4098a4c6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 0}]\n",
              "                    [{TPU 1}]], shape=(2, 1))"
            ]
          },
          "execution_count": 79,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sharding_a"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7lmmW2UkaFqY",
        "outputId": "9b67c75f-c363-4a2e-de3e-31450276b59e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "PositionalSharding([[{TPU 2}]\n",
              "                    [{TPU 3}]], shape=(2, 1))"
            ]
          },
          "execution_count": 80,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sharding_b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "id": "VQ-IkTSmaKst"
      },
      "outputs": [],
      "source": [
        "v1sp = jax.device_put(v1s, sharding_a)\n",
        "v2sp = jax.device_put(v2s, sharding_b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 82,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "cLk4ddMBboyI",
        "outputId": "01b27140-aee2-41fa-cedc-02b9d4f61bd3"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 1                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "id": "88sPq_RZbpnY",
        "outputId": "5504d57a-8218-4bdf-8f16-e89a03e9b126"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 2                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 3                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 453
        },
        "id": "WYBV-u87aNs1",
        "outputId": "352a452e-244e-49a2-ec9e-28e2b2dce746"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1] on platform TPU and ARG_SHARDING with device ids [2, 3] on platform TPU",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[84], 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[7], 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~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/pjit.py:197\u001b[0m, in \u001b[0;36m_python_pjit_helper\u001b[0;34m(fun, jit_info, *args, **kwargs)\u001b[0m\n\u001b[1;32m    194\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    195\u001b[0m   msg \u001b[38;5;241m=\u001b[39m _device_assignment_mismatch_error(\n\u001b[1;32m    196\u001b[0m       fun_name, fails, args_flat, api_name, p\u001b[38;5;241m.\u001b[39marg_names)\n\u001b[0;32m--> 197\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    198\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m xla\u001b[38;5;241m.\u001b[39mInvalidInputException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    199\u001b[0m   arg_names \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(args_flat) \u001b[38;5;28;01mif\u001b[39;00m p\u001b[38;5;241m.\u001b[39marg_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m p\u001b[38;5;241m.\u001b[39marg_names\n",
            "\u001b[0;31mValueError\u001b[0m: Received incompatible devices for jitted computation. Got ARG_SHARDING with device ids [0, 1] on platform TPU and ARG_SHARDING with device ids [2, 3] on platform TPU"
          ]
        }
      ],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0C5PKKelbP8n"
      },
      "source": [
        "Different order:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 85,
      "metadata": {
        "id": "XV6jtqLSdKqC"
      },
      "outputs": [],
      "source": [
        "sharding_a = PositionalSharding(np.array(jax.devices()).reshape(4,1))\n",
        "sharding_b = PositionalSharding(np.array(jax.devices()[::-1]).reshape(4,1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-2ly9uCNdODU",
        "outputId": "23d8b552-2453-4235-ebb3-a8bbf80e6630"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(PositionalSharding([[{TPU 0}]\n",
              "                     [{TPU 1}]\n",
              "                     [{TPU 2}]\n",
              "                     [{TPU 3}]], shape=(4, 1)),\n",
              " PositionalSharding([[{TPU 3}]\n",
              "                     [{TPU 2}]\n",
              "                     [{TPU 1}]\n",
              "                     [{TPU 0}]], shape=(4, 1)))"
            ]
          },
          "execution_count": 86,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sharding_a, sharding_b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 87,
      "metadata": {
        "id": "jW8JqD73dW0u"
      },
      "outputs": [],
      "source": [
        "v1sp = jax.device_put(v1s, sharding_a)\n",
        "v2sp = jax.device_put(v2s, sharding_b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "5MM-Gd4fdXL4",
        "outputId": "6b1d87be-acb9-4db3-fa96-4e035dc77a0f"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 0                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                            TPU 1                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                            TPU 2                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                            TPU 3                             </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                            \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                             \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v1sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 301
        },
        "id": "Vzm3d9OWdZ7b",
        "outputId": "4b31a0c5-a221-498a-f6cf-81704f447649"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                            TPU 3                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                            TPU 2                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                            TPU 1                             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                            TPU 0                             </span>\n",
              "<span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">                                                              </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                            \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;222;158;214m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                            \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 0\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                             \u001b[0m\n",
              "\u001b[38;2;0;0;0;48;2;181;207;107m                                                              \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 453
        },
        "id": "JsIjQdCtdeIg",
        "outputId": "e51aadfd-e770-4a3d-81c8-f854f66ba604"
      },
      "outputs": [
        {
          "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 [3, 2, 1, 0] on platform TPU",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[90], 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[7], 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~/JAX-in-Action/.venv/lib/python3.10/site-packages/jax/_src/pjit.py:197\u001b[0m, in \u001b[0;36m_python_pjit_helper\u001b[0;34m(fun, jit_info, *args, **kwargs)\u001b[0m\n\u001b[1;32m    194\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    195\u001b[0m   msg \u001b[38;5;241m=\u001b[39m _device_assignment_mismatch_error(\n\u001b[1;32m    196\u001b[0m       fun_name, fails, args_flat, api_name, p\u001b[38;5;241m.\u001b[39marg_names)\n\u001b[0;32m--> 197\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    198\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m xla\u001b[38;5;241m.\u001b[39mInvalidInputException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    199\u001b[0m   arg_names \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mlen\u001b[39m(args_flat) \u001b[38;5;28;01mif\u001b[39;00m p\u001b[38;5;241m.\u001b[39marg_names \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m p\u001b[38;5;241m.\u001b[39marg_names\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 [3, 2, 1, 0] on platform TPU"
          ]
        }
      ],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2sp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "id": "_Z7Mlk2KfPW3"
      },
      "outputs": [],
      "source": [
        "d = jax.vmap(dot)(v1sp, v2s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "id": "FdcPLXpPfQrc",
        "outputId": "74277cd0-92a0-4c6f-d3c7-acdc5804bb2f"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(d)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uwkURlR91YHR"
      },
      "source": [
        "## MLP example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NZySdFPR1fWr"
      },
      "source": [
        "### Preparing data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wFW5RFD2NqOI"
      },
      "source": [
        "Install these modules if you created a new empty cloud machine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lvgUueADKium",
        "outputId": "cd9a364e-17d2-499e-e01d-67d100c1d7ed"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R1PpLK59K12_",
        "outputId": "8189efb7-51ce-4e29-d5ab-ff91d4952a54"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BrWQqOqTuz8C",
        "outputId": "6e0718e6-b014-4544-d229-8c2890d8c179"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-12 04:25:43.085096: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
            "2024-08-12 04:25:43.107179: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
            "2024-08-12 04:25:43.113934: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
            "2024-08-12 04:25:44.060678: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\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": 94,
      "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": 95,
      "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": 96,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "seqmmqqx1let",
        "outputId": "1df1ebe1-264f-4518-da9f-4656104d4299"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "469"
            ]
          },
          "execution_count": 96,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TkQT5al_1ove"
      },
      "source": [
        "### Preparing MLP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "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": 98,
      "metadata": {
        "id": "oRrZN0Xh1qsl"
      },
      "outputs": [],
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "PARAM_SCALE = 0.01"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "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": 100,
      "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": 101,
      "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": 102,
      "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 - jnp.expand_dims(logsumexp(logits, axis=1), 1) # logsumexp trick https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/\n",
        "  return -jnp.mean(targets*log_preds)\n",
        "\n",
        "@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": 103,
      "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",
      "metadata": {
        "id": "hzF2I-4V3Qjv"
      },
      "source": [
        "#### 8-way data parallelism"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {
        "id": "JrhlyerwzgZ1"
      },
      "outputs": [],
      "source": [
        "sharding = PositionalSharding(jax.devices()).reshape(4, 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 109,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CGd4pczh2I1h",
        "outputId": "8a870220-484d-4095-ceb2-1f0402a7daa9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 0 in 0.92 sec\n",
            "Training set loss 0.0822610855102539\n",
            "Training set accuracy 0.8977934122085571\n",
            "Test set accuracy 0.9017009735107422\n",
            "Epoch 1 in 0.88 sec\n",
            "Training set loss 0.033943504095077515\n",
            "Training set accuracy 0.9136738777160645\n",
            "Test set accuracy 0.915249228477478\n",
            "Epoch 2 in 0.97 sec\n",
            "Training set loss 0.030296433717012405\n",
            "Training set accuracy 0.9209200143814087\n",
            "Test set accuracy 0.9210838675498962\n",
            "Epoch 3 in 0.96 sec\n",
            "Training set loss 0.027934584766626358\n",
            "Training set accuracy 0.9271888732910156\n",
            "Test set accuracy 0.9252373576164246\n",
            "Epoch 4 in 0.97 sec\n",
            "Training set loss 0.025554420426487923\n",
            "Training set accuracy 0.9339630603790283\n",
            "Test set accuracy 0.9323576092720032\n",
            "Epoch 5 in 0.99 sec\n",
            "Training set loss 0.023182667791843414\n",
            "Training set accuracy 0.9401763677597046\n",
            "Test set accuracy 0.938686728477478\n",
            "Epoch 6 in 0.95 sec\n",
            "Training set loss 0.02106645330786705\n",
            "Training set accuracy 0.9453791379928589\n",
            "Test set accuracy 0.9431368708610535\n",
            "Epoch 7 in 0.92 sec\n",
            "Training set loss 0.019250400364398956\n",
            "Training set accuracy 0.9495935440063477\n",
            "Test set accuracy 0.9471914768218994\n",
            "Epoch 8 in 1.01 sec\n",
            "Training set loss 0.017694123089313507\n",
            "Training set accuracy 0.9541911482810974\n",
            "Test set accuracy 0.9508504867553711\n",
            "Epoch 9 in 0.95 sec\n",
            "Training set loss 0.016348740085959435\n",
            "Training set accuracy 0.9573394656181335\n",
            "Test set accuracy 0.953718364238739\n",
            "Epoch 10 in 0.91 sec\n",
            "Training set loss 0.015177320688962936\n",
            "Training set accuracy 0.9601046442985535\n",
            "Test set accuracy 0.9567840695381165\n",
            "Epoch 11 in 0.86 sec\n",
            "Training set loss 0.014149891212582588\n",
            "Training set accuracy 0.9625700116157532\n",
            "Test set accuracy 0.9581685662269592\n",
            "Epoch 12 in 0.94 sec\n",
            "Training set loss 0.013246104121208191\n",
            "Training set accuracy 0.9648021459579468\n",
            "Test set accuracy 0.9605419635772705\n",
            "Epoch 13 in 0.91 sec\n",
            "Training set loss 0.012447650544345379\n",
            "Training set accuracy 0.9667844176292419\n",
            "Test set accuracy 0.9619264602661133\n",
            "Epoch 14 in 1.01 sec\n",
            "Training set loss 0.011738198809325695\n",
            "Training set accuracy 0.968533456325531\n",
            "Test set accuracy 0.9625198245048523\n",
            "Epoch 15 in 0.94 sec\n",
            "Training set loss 0.011108276434242725\n",
            "Training set accuracy 0.9700326919555664\n",
            "Test set accuracy 0.9640032052993774\n",
            "Epoch 16 in 1.21 sec\n",
            "Training set loss 0.010542847216129303\n",
            "Training set accuracy 0.9714319109916687\n",
            "Test set accuracy 0.9656843543052673\n",
            "Epoch 17 in 1.19 sec\n",
            "Training set loss 0.010033256374299526\n",
            "Training set accuracy 0.9728978276252747\n",
            "Test set accuracy 0.966178834438324\n",
            "Epoch 18 in 1.42 sec\n",
            "Training set loss 0.00957263819873333\n",
            "Training set accuracy 0.9741471409797668\n",
            "Test set accuracy 0.9673655033111572\n",
            "Epoch 19 in 1.01 sec\n",
            "Training set loss 0.009151486679911613\n",
            "Training set accuracy 0.9752465486526489\n",
            "Test set accuracy 0.96875\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(loss_value)\n",
        "  epoch_time = time.time() - start_time\n",
        "\n",
        "  train_acc = accuracy(params, train_data)\n",
        "  test_acc = accuracy(params, test_data)\n",
        "  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
        "  print(\"Training set loss {}\".format(jnp.mean(jnp.array(losses))))\n",
        "  print(\"Training set accuracy {}\".format(train_acc))\n",
        "  print(\"Test set accuracy {}\".format(test_acc))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 108,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(469,\n",
              " [Array(0.00715537, dtype=float32),\n",
              "  Array(0.00720292, dtype=float32),\n",
              "  Array(0.00704907, dtype=float32),\n",
              "  Array(0.00552397, dtype=float32),\n",
              "  Array(0.00984609, dtype=float32),\n",
              "  Array(0.01049436, dtype=float32),\n",
              "  Array(0.00601297, dtype=float32),\n",
              "  Array(0.0108844, dtype=float32),\n",
              "  Array(0.00592147, dtype=float32),\n",
              "  Array(0.00666778, dtype=float32),\n",
              "  Array(0.00456971, dtype=float32),\n",
              "  Array(0.00707005, dtype=float32),\n",
              "  Array(0.00967919, dtype=float32),\n",
              "  Array(0.01597867, dtype=float32),\n",
              "  Array(0.01205531, dtype=float32),\n",
              "  Array(0.00985306, dtype=float32),\n",
              "  Array(0.00408849, dtype=float32),\n",
              "  Array(0.01805894, dtype=float32),\n",
              "  Array(0.01227613, dtype=float32),\n",
              "  Array(0.00526123, dtype=float32),\n",
              "  Array(0.01130037, dtype=float32),\n",
              "  Array(0.01170993, dtype=float32),\n",
              "  Array(0.00646852, dtype=float32),\n",
              "  Array(0.02234733, dtype=float32),\n",
              "  Array(0.01576995, dtype=float32),\n",
              "  Array(0.00691734, dtype=float32),\n",
              "  Array(0.01684645, dtype=float32),\n",
              "  Array(0.00574006, dtype=float32),\n",
              "  Array(0.00784057, dtype=float32),\n",
              "  Array(0.01103384, dtype=float32),\n",
              "  Array(0.00618663, dtype=float32),\n",
              "  Array(0.01186787, dtype=float32),\n",
              "  Array(0.01008057, dtype=float32),\n",
              "  Array(0.00708062, dtype=float32),\n",
              "  Array(0.00904311, dtype=float32),\n",
              "  Array(0.00686354, dtype=float32),\n",
              "  Array(0.00595688, dtype=float32),\n",
              "  Array(0.01245286, dtype=float32),\n",
              "  Array(0.01101474, dtype=float32),\n",
              "  Array(0.00846246, dtype=float32),\n",
              "  Array(0.00817709, dtype=float32),\n",
              "  Array(0.01122988, dtype=float32),\n",
              "  Array(0.01849934, dtype=float32),\n",
              "  Array(0.00982656, dtype=float32),\n",
              "  Array(0.00446335, dtype=float32),\n",
              "  Array(0.01335639, dtype=float32),\n",
              "  Array(0.0034414, dtype=float32),\n",
              "  Array(0.00558565, dtype=float32),\n",
              "  Array(0.01501527, dtype=float32),\n",
              "  Array(0.00852545, dtype=float32),\n",
              "  Array(0.01511432, dtype=float32),\n",
              "  Array(0.01330939, dtype=float32),\n",
              "  Array(0.01065705, dtype=float32),\n",
              "  Array(0.00558391, dtype=float32),\n",
              "  Array(0.00831316, dtype=float32),\n",
              "  Array(0.00443693, dtype=float32),\n",
              "  Array(0.01320294, dtype=float32),\n",
              "  Array(0.00650542, dtype=float32),\n",
              "  Array(0.00595414, dtype=float32),\n",
              "  Array(0.0072539, dtype=float32),\n",
              "  Array(0.01372101, dtype=float32),\n",
              "  Array(0.01480604, dtype=float32),\n",
              "  Array(0.01027423, dtype=float32),\n",
              "  Array(0.00668659, dtype=float32),\n",
              "  Array(0.00490806, dtype=float32),\n",
              "  Array(0.01111231, dtype=float32),\n",
              "  Array(0.00642464, dtype=float32),\n",
              "  Array(0.00497197, dtype=float32),\n",
              "  Array(0.00929168, dtype=float32),\n",
              "  Array(0.00695069, dtype=float32),\n",
              "  Array(0.00619658, dtype=float32),\n",
              "  Array(0.00638385, dtype=float32),\n",
              "  Array(0.00764535, dtype=float32),\n",
              "  Array(0.00593925, dtype=float32),\n",
              "  Array(0.01405862, dtype=float32),\n",
              "  Array(0.00593139, dtype=float32),\n",
              "  Array(0.01769333, dtype=float32),\n",
              "  Array(0.00403769, dtype=float32),\n",
              "  Array(0.00698057, dtype=float32),\n",
              "  Array(0.00967532, dtype=float32),\n",
              "  Array(0.01067164, dtype=float32),\n",
              "  Array(0.01459931, dtype=float32),\n",
              "  Array(0.015528, dtype=float32),\n",
              "  Array(0.00760492, dtype=float32),\n",
              "  Array(0.01162664, dtype=float32),\n",
              "  Array(0.0097457, dtype=float32),\n",
              "  Array(0.0058027, dtype=float32),\n",
              "  Array(0.00534712, dtype=float32),\n",
              "  Array(0.00638707, dtype=float32),\n",
              "  Array(0.00374175, dtype=float32),\n",
              "  Array(0.01170266, dtype=float32),\n",
              "  Array(0.00653279, dtype=float32),\n",
              "  Array(0.02663139, dtype=float32),\n",
              "  Array(0.00496221, dtype=float32),\n",
              "  Array(0.00938921, dtype=float32),\n",
              "  Array(0.01511282, dtype=float32),\n",
              "  Array(0.00975549, dtype=float32),\n",
              "  Array(0.00753017, dtype=float32),\n",
              "  Array(0.0033096, dtype=float32),\n",
              "  Array(0.00838848, dtype=float32),\n",
              "  Array(0.0047805, dtype=float32),\n",
              "  Array(0.01049828, dtype=float32),\n",
              "  Array(0.00622318, dtype=float32),\n",
              "  Array(0.00675792, dtype=float32),\n",
              "  Array(0.00476324, dtype=float32),\n",
              "  Array(0.01110208, dtype=float32),\n",
              "  Array(0.00698493, dtype=float32),\n",
              "  Array(0.01154176, dtype=float32),\n",
              "  Array(0.0094326, dtype=float32),\n",
              "  Array(0.01253151, dtype=float32),\n",
              "  Array(0.00572668, dtype=float32),\n",
              "  Array(0.00817569, dtype=float32),\n",
              "  Array(0.00907806, dtype=float32),\n",
              "  Array(0.00617611, dtype=float32),\n",
              "  Array(0.01427079, dtype=float32),\n",
              "  Array(0.00374501, dtype=float32),\n",
              "  Array(0.00784046, dtype=float32),\n",
              "  Array(0.00966949, dtype=float32),\n",
              "  Array(0.00358428, dtype=float32),\n",
              "  Array(0.00952203, dtype=float32),\n",
              "  Array(0.01125365, dtype=float32),\n",
              "  Array(0.01078991, dtype=float32),\n",
              "  Array(0.00745513, dtype=float32),\n",
              "  Array(0.00946437, dtype=float32),\n",
              "  Array(0.00518008, dtype=float32),\n",
              "  Array(0.00950578, dtype=float32),\n",
              "  Array(0.00608064, dtype=float32),\n",
              "  Array(0.00849207, dtype=float32),\n",
              "  Array(0.01761059, dtype=float32),\n",
              "  Array(0.00671197, dtype=float32),\n",
              "  Array(0.00335123, dtype=float32),\n",
              "  Array(0.01362847, dtype=float32),\n",
              "  Array(0.00982094, dtype=float32),\n",
              "  Array(0.01410992, dtype=float32),\n",
              "  Array(0.00809181, dtype=float32),\n",
              "  Array(0.00820133, dtype=float32),\n",
              "  Array(0.0112866, dtype=float32),\n",
              "  Array(0.00825891, dtype=float32),\n",
              "  Array(0.00550814, dtype=float32),\n",
              "  Array(0.00507368, dtype=float32),\n",
              "  Array(0.00597332, dtype=float32),\n",
              "  Array(0.01366306, dtype=float32),\n",
              "  Array(0.00646757, dtype=float32),\n",
              "  Array(0.01358973, dtype=float32),\n",
              "  Array(0.00932815, dtype=float32),\n",
              "  Array(0.00374975, dtype=float32),\n",
              "  Array(0.00149788, dtype=float32),\n",
              "  Array(0.0122201, dtype=float32),\n",
              "  Array(0.00467177, dtype=float32),\n",
              "  Array(0.00556599, dtype=float32),\n",
              "  Array(0.01006958, dtype=float32),\n",
              "  Array(0.00942326, dtype=float32),\n",
              "  Array(0.00576596, dtype=float32),\n",
              "  Array(0.0068546, dtype=float32),\n",
              "  Array(0.00693593, dtype=float32),\n",
              "  Array(0.0031088, dtype=float32),\n",
              "  Array(0.00827179, dtype=float32),\n",
              "  Array(0.01162764, dtype=float32),\n",
              "  Array(0.0050324, dtype=float32),\n",
              "  Array(0.00881435, dtype=float32),\n",
              "  Array(0.01186213, dtype=float32),\n",
              "  Array(0.01131363, dtype=float32),\n",
              "  Array(0.00743537, dtype=float32),\n",
              "  Array(0.00737761, dtype=float32),\n",
              "  Array(0.00864948, dtype=float32),\n",
              "  Array(0.01767044, dtype=float32),\n",
              "  Array(0.00354798, dtype=float32),\n",
              "  Array(0.00542418, dtype=float32),\n",
              "  Array(0.00466838, dtype=float32),\n",
              "  Array(0.01158668, dtype=float32),\n",
              "  Array(0.00821351, dtype=float32),\n",
              "  Array(0.00527863, dtype=float32),\n",
              "  Array(0.00710708, dtype=float32),\n",
              "  Array(0.01405563, dtype=float32),\n",
              "  Array(0.0074969, dtype=float32),\n",
              "  Array(0.00894281, dtype=float32),\n",
              "  Array(0.00870649, dtype=float32),\n",
              "  Array(0.01554444, dtype=float32),\n",
              "  Array(0.00835336, dtype=float32),\n",
              "  Array(0.0116499, dtype=float32),\n",
              "  Array(0.00697969, dtype=float32),\n",
              "  Array(0.00963759, dtype=float32),\n",
              "  Array(0.00456813, dtype=float32),\n",
              "  Array(0.0061488, dtype=float32),\n",
              "  Array(0.00624182, dtype=float32),\n",
              "  Array(0.01009657, dtype=float32),\n",
              "  Array(0.01257215, dtype=float32),\n",
              "  Array(0.01051517, dtype=float32),\n",
              "  Array(0.00882395, dtype=float32),\n",
              "  Array(0.00517154, dtype=float32),\n",
              "  Array(0.01185858, dtype=float32),\n",
              "  Array(0.01224419, dtype=float32),\n",
              "  Array(0.00792753, dtype=float32),\n",
              "  Array(0.01117677, dtype=float32),\n",
              "  Array(0.00946368, dtype=float32),\n",
              "  Array(0.0071283, dtype=float32),\n",
              "  Array(0.01711252, dtype=float32),\n",
              "  Array(0.00692343, dtype=float32),\n",
              "  Array(0.01329454, dtype=float32),\n",
              "  Array(0.00547658, dtype=float32),\n",
              "  Array(0.0109752, dtype=float32),\n",
              "  Array(0.00858221, dtype=float32),\n",
              "  Array(0.00805087, dtype=float32),\n",
              "  Array(0.00571564, dtype=float32),\n",
              "  Array(0.00378814, dtype=float32),\n",
              "  Array(0.00337793, dtype=float32),\n",
              "  Array(0.01525816, dtype=float32),\n",
              "  Array(0.00757033, dtype=float32),\n",
              "  Array(0.00992832, dtype=float32),\n",
              "  Array(0.01096067, dtype=float32),\n",
              "  Array(0.01809762, dtype=float32),\n",
              "  Array(0.00938112, dtype=float32),\n",
              "  Array(0.01062687, dtype=float32),\n",
              "  Array(0.0154736, dtype=float32),\n",
              "  Array(0.01373984, dtype=float32),\n",
              "  Array(0.01492539, dtype=float32),\n",
              "  Array(0.0088467, dtype=float32),\n",
              "  Array(0.00408829, dtype=float32),\n",
              "  Array(0.00941716, dtype=float32),\n",
              "  Array(0.00386626, dtype=float32),\n",
              "  Array(0.01372675, dtype=float32),\n",
              "  Array(0.00503443, dtype=float32),\n",
              "  Array(0.00952703, dtype=float32),\n",
              "  Array(0.00354555, dtype=float32),\n",
              "  Array(0.0055552, dtype=float32),\n",
              "  Array(0.00731229, dtype=float32),\n",
              "  Array(0.00930953, dtype=float32),\n",
              "  Array(0.01111909, dtype=float32),\n",
              "  Array(0.0068629, dtype=float32),\n",
              "  Array(0.00934988, dtype=float32),\n",
              "  Array(0.00634522, dtype=float32),\n",
              "  Array(0.00617057, dtype=float32),\n",
              "  Array(0.00448716, dtype=float32),\n",
              "  Array(0.00847401, dtype=float32),\n",
              "  Array(0.01900866, dtype=float32),\n",
              "  Array(0.01450234, dtype=float32),\n",
              "  Array(0.00981516, dtype=float32),\n",
              "  Array(0.00535012, dtype=float32),\n",
              "  Array(0.0120869, dtype=float32),\n",
              "  Array(0.00353178, dtype=float32),\n",
              "  Array(0.01394098, dtype=float32),\n",
              "  Array(0.00860516, dtype=float32),\n",
              "  Array(0.01613839, dtype=float32),\n",
              "  Array(0.01125447, dtype=float32),\n",
              "  Array(0.01349407, dtype=float32),\n",
              "  Array(0.002623, dtype=float32),\n",
              "  Array(0.00480208, dtype=float32),\n",
              "  Array(0.00400821, dtype=float32),\n",
              "  Array(0.0125119, dtype=float32),\n",
              "  Array(0.00576411, dtype=float32),\n",
              "  Array(0.00605031, dtype=float32),\n",
              "  Array(0.00823624, dtype=float32),\n",
              "  Array(0.00777326, dtype=float32),\n",
              "  Array(0.00411577, dtype=float32),\n",
              "  Array(0.00823261, dtype=float32),\n",
              "  Array(0.00198878, dtype=float32),\n",
              "  Array(0.0066186, dtype=float32),\n",
              "  Array(0.00539755, dtype=float32),\n",
              "  Array(0.00498512, dtype=float32),\n",
              "  Array(0.00589674, dtype=float32),\n",
              "  Array(0.0071999, dtype=float32),\n",
              "  Array(0.02529567, dtype=float32),\n",
              "  Array(0.01442621, dtype=float32),\n",
              "  Array(0.01427113, dtype=float32),\n",
              "  Array(0.01104017, dtype=float32),\n",
              "  Array(0.0055207, dtype=float32),\n",
              "  Array(0.00668925, dtype=float32),\n",
              "  Array(0.01019346, dtype=float32),\n",
              "  Array(0.01188943, dtype=float32),\n",
              "  Array(0.00395488, dtype=float32),\n",
              "  Array(0.01033677, dtype=float32),\n",
              "  Array(0.00956674, dtype=float32),\n",
              "  Array(0.00854938, dtype=float32),\n",
              "  Array(0.00884258, dtype=float32),\n",
              "  Array(0.00995028, dtype=float32),\n",
              "  Array(0.00535104, dtype=float32),\n",
              "  Array(0.00952279, dtype=float32),\n",
              "  Array(0.01593888, dtype=float32),\n",
              "  Array(0.01268555, dtype=float32),\n",
              "  Array(0.01175113, dtype=float32),\n",
              "  Array(0.0085984, dtype=float32),\n",
              "  Array(0.00520565, dtype=float32),\n",
              "  Array(0.00456704, dtype=float32),\n",
              "  Array(0.01502567, dtype=float32),\n",
              "  Array(0.00751545, dtype=float32),\n",
              "  Array(0.00626393, dtype=float32),\n",
              "  Array(0.0063555, dtype=float32),\n",
              "  Array(0.01148976, dtype=float32),\n",
              "  Array(0.00952704, dtype=float32),\n",
              "  Array(0.00882051, dtype=float32),\n",
              "  Array(0.00748168, dtype=float32),\n",
              "  Array(0.00947486, dtype=float32),\n",
              "  Array(0.01203895, dtype=float32),\n",
              "  Array(0.00721907, dtype=float32),\n",
              "  Array(0.01415384, dtype=float32),\n",
              "  Array(0.01706646, dtype=float32),\n",
              "  Array(0.00620816, dtype=float32),\n",
              "  Array(0.00854848, dtype=float32),\n",
              "  Array(0.0099365, dtype=float32),\n",
              "  Array(0.00610333, dtype=float32),\n",
              "  Array(0.01284967, dtype=float32),\n",
              "  Array(0.00503805, dtype=float32),\n",
              "  Array(0.00794949, dtype=float32),\n",
              "  Array(0.01163037, dtype=float32),\n",
              "  Array(0.01075088, dtype=float32),\n",
              "  Array(0.00920765, dtype=float32),\n",
              "  Array(0.00904822, dtype=float32),\n",
              "  Array(0.01816061, dtype=float32),\n",
              "  Array(0.00762138, dtype=float32),\n",
              "  Array(0.00747412, dtype=float32),\n",
              "  Array(0.00809723, dtype=float32),\n",
              "  Array(0.00484933, dtype=float32),\n",
              "  Array(0.01330119, dtype=float32),\n",
              "  Array(0.01228002, dtype=float32),\n",
              "  Array(0.00569284, dtype=float32),\n",
              "  Array(0.01129412, dtype=float32),\n",
              "  Array(0.00744513, dtype=float32),\n",
              "  Array(0.00578101, dtype=float32),\n",
              "  Array(0.00469105, dtype=float32),\n",
              "  Array(0.003395, dtype=float32),\n",
              "  Array(0.00617441, dtype=float32),\n",
              "  Array(0.02023105, dtype=float32),\n",
              "  Array(0.00684483, dtype=float32),\n",
              "  Array(0.0127619, dtype=float32),\n",
              "  Array(0.00862186, dtype=float32),\n",
              "  Array(0.01015411, dtype=float32),\n",
              "  Array(0.00355815, dtype=float32),\n",
              "  Array(0.00398131, dtype=float32),\n",
              "  Array(0.00939332, dtype=float32),\n",
              "  Array(0.00993729, dtype=float32),\n",
              "  Array(0.00702658, dtype=float32),\n",
              "  Array(0.01195145, dtype=float32),\n",
              "  Array(0.00849271, dtype=float32),\n",
              "  Array(0.01472493, dtype=float32),\n",
              "  Array(0.00537741, dtype=float32),\n",
              "  Array(0.00631062, dtype=float32),\n",
              "  Array(0.00849116, dtype=float32),\n",
              "  Array(0.00666478, dtype=float32),\n",
              "  Array(0.00840349, dtype=float32),\n",
              "  Array(0.00822061, dtype=float32),\n",
              "  Array(0.00435342, dtype=float32),\n",
              "  Array(0.00607981, dtype=float32),\n",
              "  Array(0.00904926, dtype=float32),\n",
              "  Array(0.00504346, dtype=float32),\n",
              "  Array(0.00736488, dtype=float32),\n",
              "  Array(0.01243273, dtype=float32),\n",
              "  Array(0.00871171, dtype=float32),\n",
              "  Array(0.00432171, dtype=float32),\n",
              "  Array(0.00810324, dtype=float32),\n",
              "  Array(0.01099622, dtype=float32),\n",
              "  Array(0.00590551, dtype=float32),\n",
              "  Array(0.01418192, dtype=float32),\n",
              "  Array(0.01505364, dtype=float32),\n",
              "  Array(0.02631317, dtype=float32),\n",
              "  Array(0.01132791, dtype=float32),\n",
              "  Array(0.00614359, dtype=float32),\n",
              "  Array(0.01683249, dtype=float32),\n",
              "  Array(0.01033848, dtype=float32),\n",
              "  Array(0.01706284, dtype=float32),\n",
              "  Array(0.00859671, dtype=float32),\n",
              "  Array(0.00717983, dtype=float32),\n",
              "  Array(0.02332587, dtype=float32),\n",
              "  Array(0.01043794, dtype=float32),\n",
              "  Array(0.00941162, dtype=float32),\n",
              "  Array(0.00940573, dtype=float32),\n",
              "  Array(0.00897963, dtype=float32),\n",
              "  Array(0.00598614, dtype=float32),\n",
              "  Array(0.00538664, dtype=float32),\n",
              "  Array(0.01086632, dtype=float32),\n",
              "  Array(0.01385744, dtype=float32),\n",
              "  Array(0.01962443, dtype=float32),\n",
              "  Array(0.00484279, dtype=float32),\n",
              "  Array(0.00729563, dtype=float32),\n",
              "  Array(0.01366792, dtype=float32),\n",
              "  Array(0.00903656, dtype=float32),\n",
              "  Array(0.01092833, dtype=float32),\n",
              "  Array(0.00599955, dtype=float32),\n",
              "  Array(0.01402906, dtype=float32),\n",
              "  Array(0.00772522, dtype=float32),\n",
              "  Array(0.00952419, dtype=float32),\n",
              "  Array(0.00482681, dtype=float32),\n",
              "  Array(0.00552027, dtype=float32),\n",
              "  Array(0.01197038, dtype=float32),\n",
              "  Array(0.00843517, dtype=float32),\n",
              "  Array(0.00965284, dtype=float32),\n",
              "  Array(0.01118127, dtype=float32),\n",
              "  Array(0.01109344, dtype=float32),\n",
              "  Array(0.01068447, dtype=float32),\n",
              "  Array(0.00944095, dtype=float32),\n",
              "  Array(0.00540696, dtype=float32),\n",
              "  Array(0.00580311, dtype=float32),\n",
              "  Array(0.00913881, dtype=float32),\n",
              "  Array(0.00732108, dtype=float32),\n",
              "  Array(0.00966205, dtype=float32),\n",
              "  Array(0.00782091, dtype=float32),\n",
              "  Array(0.01202113, dtype=float32),\n",
              "  Array(0.0079737, dtype=float32),\n",
              "  Array(0.00834753, dtype=float32),\n",
              "  Array(0.01065987, dtype=float32),\n",
              "  Array(0.01009172, dtype=float32),\n",
              "  Array(0.01334727, dtype=float32),\n",
              "  Array(0.00619839, dtype=float32),\n",
              "  Array(0.00935352, dtype=float32),\n",
              "  Array(0.01164689, dtype=float32),\n",
              "  Array(0.01005257, dtype=float32),\n",
              "  Array(0.00444788, dtype=float32),\n",
              "  Array(0.00687535, dtype=float32),\n",
              "  Array(0.00864482, dtype=float32),\n",
              "  Array(0.00799138, dtype=float32),\n",
              "  Array(0.00804868, dtype=float32),\n",
              "  Array(0.01083802, dtype=float32),\n",
              "  Array(0.0127744, dtype=float32),\n",
              "  Array(0.0313906, dtype=float32),\n",
              "  Array(0.00615665, dtype=float32),\n",
              "  Array(0.01300749, dtype=float32),\n",
              "  Array(0.00917291, dtype=float32),\n",
              "  Array(0.00849073, dtype=float32),\n",
              "  Array(0.00621255, dtype=float32),\n",
              "  Array(0.00772546, dtype=float32),\n",
              "  Array(0.0108934, dtype=float32),\n",
              "  Array(0.00513792, dtype=float32),\n",
              "  Array(0.01235466, dtype=float32),\n",
              "  Array(0.00491095, dtype=float32),\n",
              "  Array(0.01030413, dtype=float32),\n",
              "  Array(0.00528422, dtype=float32),\n",
              "  Array(0.00654849, dtype=float32),\n",
              "  Array(0.01158838, dtype=float32),\n",
              "  Array(0.00991549, dtype=float32),\n",
              "  Array(0.01568046, dtype=float32),\n",
              "  Array(0.01179941, dtype=float32),\n",
              "  Array(0.01269739, dtype=float32),\n",
              "  Array(0.00507445, dtype=float32),\n",
              "  Array(0.00334293, dtype=float32),\n",
              "  Array(0.00362204, dtype=float32),\n",
              "  Array(0.01183572, dtype=float32),\n",
              "  Array(0.00653935, dtype=float32),\n",
              "  Array(0.00902529, dtype=float32),\n",
              "  Array(0.01116293, dtype=float32),\n",
              "  Array(0.00909121, dtype=float32),\n",
              "  Array(0.00605366, dtype=float32),\n",
              "  Array(0.01107628, dtype=float32),\n",
              "  Array(0.00920255, dtype=float32),\n",
              "  Array(0.0049389, dtype=float32),\n",
              "  Array(0.01118881, dtype=float32),\n",
              "  Array(0.0064004, dtype=float32),\n",
              "  Array(0.00442893, dtype=float32),\n",
              "  Array(0.00947342, dtype=float32),\n",
              "  Array(0.0050448, dtype=float32),\n",
              "  Array(0.01102102, dtype=float32),\n",
              "  Array(0.00573556, dtype=float32),\n",
              "  Array(0.00918462, dtype=float32),\n",
              "  Array(0.0067548, dtype=float32),\n",
              "  Array(0.00958075, dtype=float32),\n",
              "  Array(0.0201444, dtype=float32),\n",
              "  Array(0.00589464, dtype=float32),\n",
              "  Array(0.01659547, dtype=float32),\n",
              "  Array(0.00325762, dtype=float32),\n",
              "  Array(0.01118485, dtype=float32),\n",
              "  Array(0.01324637, dtype=float32),\n",
              "  Array(0.00550017, dtype=float32),\n",
              "  Array(0.00685694, dtype=float32),\n",
              "  Array(0.00944922, dtype=float32),\n",
              "  Array(0.0042129, dtype=float32),\n",
              "  Array(0.00581109, dtype=float32),\n",
              "  Array(0.020259, dtype=float32),\n",
              "  Array(0.00307948, dtype=float32),\n",
              "  Array(0.01733901, dtype=float32),\n",
              "  Array(0.00884266, dtype=float32),\n",
              "  Array(0.00335006, dtype=float32)])"
            ]
          },
          "execution_count": 108,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "\n",
        "len(losses), losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 110,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "id": "FZxrQ5OvhV7u",
        "outputId": "5457df85-c04d-4412-d132-0a55ca2626cc"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">             TPU 0,1,2,3             </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m             \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "jax.debug.visualize_array_sharding(params[0][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alhuQ3jN3Zp8"
      },
      "source": [
        "#### 4-way data parallelism, 2-way tensor parallelism"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 162,
      "metadata": {
        "id": "vLmoPdbu3Y8p"
      },
      "outputs": [],
      "source": [
        "sharding = PositionalSharding(jax.devices()).reshape(2, 2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 163,
      "metadata": {
        "id": "jyVcAxPC3mGG"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 164,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-5dPxHms4a5E",
        "outputId": "00f41938-9fa8-4794-b679-7ad6aaf9699c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0 (10000, 784) (10000,)\n",
            "1 (10000, 10000) (10000,)\n",
            "2 (10, 10000) (10,)\n"
          ]
        }
      ],
      "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())\n",
        "  elif i==1:\n",
        "    w = jax.device_put(w, sharding.replicate(0))\n",
        "    b = jax.device_put(b, sharding.replicate())\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 165,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "1mmndGm66zkx",
        "outputId": "7480949b-83d7-45ec-e23a-d9e87b1c49ee"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">          TPU 0          </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m          \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m          \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                     TPU 0                                      </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                      \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "for (w,b) in init_params:\n",
        "  jax.debug.visualize_array_sharding(w)\n",
        "  jax.debug.visualize_array_sharding(b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 166,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "vJaVTiEP8V43",
        "outputId": "1a4b282d-f803-4e7f-84f6-49128c5c2b06"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">       TPU 0,1,2,3       </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m       \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m       \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  TPU 0,1,2,3                                   </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "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()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 167,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 720
        },
        "id": "o28ff-t966yS",
        "outputId": "dba7c757-2054-4915-e904-03f93d134ce2"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">   TPU 0    </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">   TPU 1    </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">            </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">            </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">            </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">   TPU 2    </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">   TPU 3    </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">            </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">            </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">            </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">            </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">            </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m   \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m    \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m   \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m    \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m            \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m   \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m    \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m   \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m    \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m            \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m            \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m            \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                  TPU 0,1,2,3                                   </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">                                                                                </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                   \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">TPU 0,1,2,3</span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">           </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0,1,2,3\u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "for (w,b) in sharded_params:\n",
        "  jax.debug.visualize_array_sharding(w)\n",
        "  jax.debug.visualize_array_sharding(b)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 168,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "9y4hEhkp7GQI",
        "outputId": "229e456d-c65e-4f72-c9db-fd3ddaf41de6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 0 in 1.01 sec\n",
            "Training set loss 0.06498625129461288\n",
            "Training set accuracy 0.8998367786407471\n",
            "Test set accuracy 0.9048655033111572\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1 in 1.00 sec\n",
            "Training set loss 0.03310569375753403\n",
            "Training set accuracy 0.9135960936546326\n",
            "Test set accuracy 0.9148536324501038\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 2 in 1.01 sec\n",
            "Training set loss 0.030245665460824966\n",
            "Training set accuracy 0.9188210964202881\n",
            "Test set accuracy 0.918710470199585\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 3 in 0.94 sec\n",
            "Training set loss 0.028722697868943214\n",
            "Training set accuracy 0.9227911829948425\n",
            "Test set accuracy 0.9221717119216919\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 4 in 0.98 sec\n",
            "Training set loss 0.027605729177594185\n",
            "Training set accuracy 0.9255896806716919\n",
            "Test set accuracy 0.9243473410606384\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 5 in 1.05 sec\n",
            "Training set loss 0.026645950973033905\n",
            "Training set accuracy 0.9283826947212219\n",
            "Test set accuracy 0.9260284900665283\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 6 in 0.92 sec\n",
            "Training set loss 0.02573257125914097\n",
            "Training set accuracy 0.930936872959137\n",
            "Test set accuracy 0.9283030033111572\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 7 in 1.02 sec\n",
            "Training set loss 0.024794965982437134\n",
            "Training set accuracy 0.9339630603790283\n",
            "Test set accuracy 0.9300830960273743\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 8 in 0.97 sec\n",
            "Training set loss 0.023775801062583923\n",
            "Training set accuracy 0.9370447397232056\n",
            "Test set accuracy 0.9322587251663208\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 9 in 0.99 sec\n",
            "Training set loss 0.02263575978577137\n",
            "Training set accuracy 0.9405428767204285\n",
            "Test set accuracy 0.9359177350997925\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 10 in 1.00 sec\n",
            "Training set loss 0.021378550678491592\n",
            "Training set accuracy 0.9442464113235474\n",
            "Test set accuracy 0.9404668211936951\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 11 in 1.02 sec\n",
            "Training set loss 0.02007247321307659\n",
            "Training set accuracy 0.9480610489845276\n",
            "Test set accuracy 0.9432357549667358\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 12 in 1.00 sec\n",
            "Training set loss 0.01881973072886467\n",
            "Training set accuracy 0.9508928656578064\n",
            "Test set accuracy 0.9471914768218994\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 13 in 0.93 sec\n",
            "Training set loss 0.01768065057694912\n",
            "Training set accuracy 0.9535414576530457\n",
            "Test set accuracy 0.9497627019882202\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 14 in 0.95 sec\n",
            "Training set loss 0.016660500317811966\n",
            "Training set accuracy 0.9561067819595337\n",
            "Test set accuracy 0.9510482549667358\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 15 in 1.09 sec\n",
            "Training set loss 0.015744229778647423\n",
            "Training set accuracy 0.9581390023231506\n",
            "Test set accuracy 0.9555973410606384\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 16 in 0.94 sec\n",
            "Training set loss 0.014912167564034462\n",
            "Training set accuracy 0.9604044556617737\n",
            "Test set accuracy 0.9575752019882202\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 17 in 0.94 sec\n",
            "Training set loss 0.014147961512207985\n",
            "Training set accuracy 0.9620869159698486\n",
            "Test set accuracy 0.9588607549667358\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 18 in 1.00 sec\n",
            "Training set loss 0.013438239693641663\n",
            "Training set accuracy 0.9641191363334656\n",
            "Test set accuracy 0.9602452516555786\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 19 in 1.02 sec\n",
            "Training set loss 0.012776298448443413\n",
            "Training set accuracy 0.9654517769813538\n",
            "Test set accuracy 0.9611353278160095\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">  TPU 0  </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">  TPU 1  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #393b79\">         </span><span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #de9ed6\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">  TPU 2  </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">  TPU 3  </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "<span style=\"color: #ffffff; text-decoration-color: #ffffff; background-color: #ad494a\">         </span><span style=\"color: #000000; text-decoration-color: #000000; background-color: #b5cf6b\">         </span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mTPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mTPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;57;59;121m         \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mTPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mTPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m  \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n",
              "\u001b[38;2;255;255;255;48;2;173;73;74m         \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m         \u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "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])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJeoJJNGMrxA"
      },
      "source": [
        "**!!! Do not forget to shutdown your Cloud TPU, or you'll spend much money on it!!!**"
      ]
    },
    {
      "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": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}