Appendix-D/JAX_in_Action_Appendix_D_xmap.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## ⚖️ Choose A or B:"
      ],
      "metadata": {
        "id": "JCirjpcFSPan"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A: Emulating multi-device system on CPU\n",
        "\n",
        "Use this section to initialize a set of virtual devices on CPU if you have no access to a multi-device system.\n",
        "\n",
        "It can also help you prototype, debug and test your multi-device code locally before running it on the expensive system.\n",
        "\n",
        "Even in the case of using Google Colab it can help you prototype faster because a CPU runtime is faster to restart."
      ],
      "metadata": {
        "id": "N5lNqNEGtLyO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'"
      ],
      "metadata": {
        "id": "ypaH8OgftR_H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "n8ip1VKgterO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Tlh75iWutfVT",
        "outputId": "72bcd8e8-5b0f-43ce-cd39-41ea331ce7cd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[CpuDevice(id=0),\n",
              " CpuDevice(id=1),\n",
              " CpuDevice(id=2),\n",
              " CpuDevice(id=3),\n",
              " CpuDevice(id=4),\n",
              " CpuDevice(id=5),\n",
              " CpuDevice(id=6),\n",
              " CpuDevice(id=7)]"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## B: Setting up TPU"
      ],
      "metadata": {
        "id": "Rmk_qf6mKCZx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make preparations according to the Appendix C or Chapter 3 example"
      ],
      "metadata": {
        "id": "RMq70OC9wVHQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ],
      "metadata": {
        "id": "q9_o0MwOu94x",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3d943fa2-760d-42b3-d1a4-120d731206dd"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in links: https://storage.googleapis.com/jax-releases/libtpu_releases.html\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\n",
            "Collecting jax[tpu]\n",
            "  Downloading jax-0.4.13.tar.gz (1.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.3 MB 5.0 MB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l-\b \b\\\b \b|\b \b/\b \b-\b \bdone\n",
            "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25l-\b \b\\\b \b|\b \bdone\n",
            "\u001b[?25h    Preparing wheel metadata ... \u001b[?25l-\b \b\\\b \b|\b \b/\b \bdone\n",
            "\u001b[?25hCollecting ml-dtypes>=0.1.0\n",
            "  Downloading ml_dtypes-0.2.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.0 MB 92.7 MB/s \n",
            "\u001b[?25hCollecting numpy>=1.21\n",
            "  Downloading numpy-1.24.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 17.3 MB 97.6 MB/s \n",
            "\u001b[?25hCollecting opt-einsum\n",
            "  Downloading opt_einsum-3.3.0-py3-none-any.whl (65 kB)\n",
            "\u001b[K     |████████████████████████████████| 65 kB 5.6 MB/s \n",
            "\u001b[?25hCollecting scipy>=1.7\n",
            "  Downloading scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (34.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 34.5 MB 99.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata>=4.6; python_version < \"3.10\" in ./.local/lib/python3.8/site-packages (from jax[tpu]) (6.8.0)\n",
            "Collecting jaxlib==0.4.13; extra == \"tpu\"\n",
            "  Downloading jaxlib-0.4.13-cp38-cp38-manylinux2014_x86_64.whl (71.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 71.6 MB 51 kB/s \n",
            "\u001b[?25hCollecting libtpu-nightly==0.1.dev20230622; extra == \"tpu\"\n",
            "  Downloading https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20230622-py3-none-any.whl (171.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 171.7 MB 41 kB/s \n",
            "\u001b[?25hRequirement already satisfied: zipp>=0.5 in /usr/lib/python3/dist-packages (from importlib-metadata>=4.6; python_version < \"3.10\"->jax[tpu]) (1.0.0)\n",
            "Building wheels for collected packages: jax\n",
            "  Building wheel for jax (PEP 517) ... \u001b[?25l-\b \b\\\b \b|\b \b/\b \b-\b \b\\\b \b|\b \b/\b \b-\b \bdone\n",
            "\u001b[?25h  Created wheel for jax: filename=jax-0.4.13-py3-none-any.whl size=1518704 sha256=012d9a17c1364415da7f6c7d76f4230b5898c900c559406a5de2ee6c9f7bbb9e\n",
            "  Stored in directory: /home/grigo/.cache/pip/wheels/46/d9/15/d2800d4089dc4c77299ac7513c6aa1036f5491edbd2bf6ba16\n",
            "Successfully built jax\n",
            "Installing collected packages: numpy, ml-dtypes, opt-einsum, scipy, jaxlib, libtpu-nightly, jax\n",
            "Successfully installed jax-0.4.13 jaxlib-0.4.13 libtpu-nightly-0.1.dev20230622 ml-dtypes-0.2.0 numpy-1.24.4 opt-einsum-3.3.0 scipy-1.10.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.lib import xla_bridge\n",
        "print(xla_bridge.get_backend().platform)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2aBy_QGmu-61",
        "outputId": "1d1bfd66-235f-4824-f089-5596f1ead0e4"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "7c2IXWjRvFbj"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.local_devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_qLo4vNRvHc-",
        "outputId": "a2683a11-5107-4cba-b442-4ebea7e6fa11"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
              " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
              " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
              " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Be sure you use JAX version >= 0.4.11"
      ],
      "metadata": {
        "id": "e1KZMzA90Fuf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jax.__version__"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zjxOeo4pz-Jc",
        "outputId": "a512d8ad-008b-4981-9cf5-e55da5790463"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.4.13'"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Introducing named axes"
      ],
      "metadata": {
        "id": "XRLNPE2Ese1D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.experimental.maps import xmap"
      ],
      "metadata": {
        "id": "e3HsZ2PIsl3t"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import random"
      ],
      "metadata": {
        "id": "-Bocy386sr3n"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Replacing vmap+pmap"
      ],
      "metadata": {
        "id": "F1wKd0cB8iu5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "2p2R6uG5tVpV"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rng_key = random.PRNGKey(42)"
      ],
      "metadata": {
        "id": "SDMWf5eIswdB"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vs = random.normal(rng_key, shape=(20_000_000,3))\n",
        "v1s = vs[:10_000_000,:].T\n",
        "v2s = vs[10_000_000:,:].T\n",
        "\n",
        "v1s.shape, v2s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "THq8XQoAs8CH",
        "outputId": "e0acc982-e1d7-492e-fceb-40d29a30168e"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 10000000), (3, 10000000))"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "v1sp = v1s.reshape((v1s.shape[0], 8, v1s.shape[1]//8))\n",
        "v2sp = v2s.reshape((v2s.shape[0], 8, v2s.shape[1]//8))\n",
        "\n",
        "v1sp.shape, v2sp.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jLM8V0BstMSD",
        "outputId": "53f1226a-abbf-4fbf-b427-66c5abf60502"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 8, 1250000), (3, 8, 1250000))"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dot_parallel = jax.pmap(\n",
        "    jax.vmap(dot, in_axes=(1,1)),\n",
        "    in_axes=(1,1)\n",
        ")"
      ],
      "metadata": {
        "id": "xf9UVhJcNHH4"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = dot_parallel(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "8rdtr2F_NI_g"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dF37XSrzNLE0",
        "outputId": "a74c91b6-c480-434c-cbbb-c6641a65aa17"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000)"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_pmap = x_pmap.reshape((x_pmap.shape[0]*x_pmap.shape[1]))\n",
        "x_pmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "spEI1WDaNPXs",
        "outputId": "13db4bd4-1103-4893-8af7-a83b46ad2cfd"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "f = xmap(dot,\n",
        "         in_axes=(\n",
        "             {1:'device', 2:'batch'},\n",
        "             {1:'device', 2:'batch'}\n",
        "         ),\n",
        "         out_axes=['device', 'batch', ...]\n",
        ")"
      ],
      "metadata": {
        "id": "ARSVrRjtu0si"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap=f(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "A9sw0MMEuhET"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wbWLvlWrurqH",
        "outputId": "78518a10-101e-415d-8ba1-2f8a70f58c2b"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(8, 1250000)"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap = x_xmap.reshape((x_xmap.shape[0]*x_xmap.shape[1]))\n",
        "x_xmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RKDLzPVlIxgU",
        "outputId": "130ae5cc-5eda-49c3-e8ca-9a13b4e2b0ec"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.all(x_xmap == x_pmap)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D1q7GU6JIRYb",
        "outputId": "5beb43a2-58f0-4572-ebac-f1390a99a51c"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Einsum for comparison"
      ],
      "metadata": {
        "id": "rjXVLyPfJUjs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np"
      ],
      "metadata": {
        "id": "Fr7LYXZuLOGC"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dots = np.einsum(\"ib,ib->b\", v1s, v2s)"
      ],
      "metadata": {
        "id": "AXwuHrGRJWjK"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dots.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qLXZw8l5Jllx",
        "outputId": "f33cf196-7574-4d28-b98a-327592f4f4d2"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.all(x_xmap == dots)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hds5LT72JoYl",
        "outputId": "46bd150b-6a43-4aef-daa7-d2feb55b3234"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Changing order of output axes"
      ],
      "metadata": {
        "id": "jJH8mOtd7hn3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "f = xmap(dot,\n",
        "         in_axes=(\n",
        "             {1:'device', 2:'batch'},\n",
        "             {1:'device', 2:'batch'}\n",
        "         ),\n",
        "         out_axes=['batch', 'device', ...]\n",
        ")"
      ],
      "metadata": {
        "id": "ZCwABOzoN72M"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap=f(v1sp,v2sp)"
      ],
      "metadata": {
        "id": "d9BsiKK2UO4G"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KMmTmgnvYeOw",
        "outputId": "90321934-1abf-4aa4-bca8-3a42cb4f7792"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1250000, 8)"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using broadcasting"
      ],
      "metadata": {
        "id": "59XJ0c-c8njL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "image = random.normal(rng_key, shape=(480,640,3)) # RGB image 640x480px\n",
        "filters = random.normal(rng_key, shape=(5,3,3))   # 5 matrix filters of size 3x3"
      ],
      "metadata": {
        "id": "VE-Ny7KmYgNC"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.scipy.signal import convolve2d"
      ],
      "metadata": {
        "id": "eIjYWOuf9j6J"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def apply_filter(channel, kernel):\n",
        "  return convolve2d(channel, kernel, mode=\"same\")"
      ],
      "metadata": {
        "id": "9Zm-dYN5AA6g"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "apply_filters_to_image = xmap(apply_filter,\n",
        "         in_axes=(\n",
        "             {2:'channel'},\n",
        "             {0:'filter'}\n",
        "         ),\n",
        "         out_axes={0:'filter', 3: 'channel'}\n",
        ")"
      ],
      "metadata": {
        "id": "ej_I9TVpAbxL"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = apply_filters_to_image(image, filters)"
      ],
      "metadata": {
        "id": "OwF7aAozBEQt"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res.shape # (filters, h, w, channels)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M2iqtd__BRdH",
        "outputId": "99e64842-e5b1-42e6-9c34-85f3ba3d6b90"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, 480, 640, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using reductions"
      ],
      "metadata": {
        "id": "X5GWbod5PkjA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "f = xmap(\n",
        "     lambda x: jnp.sum(x, axis=['row']),\n",
        "     in_axes=['row', 'col'],\n",
        "     out_axes=['col']\n",
        "  )"
      ],
      "metadata": {
        "id": "ug4AF7AjAYjM"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "C = jnp.array([\n",
        "    [1,2,3],\n",
        "    [4,5,6],\n",
        "    [7,8,9]\n",
        "])"
      ],
      "metadata": {
        "id": "NkctSlJEPqdw"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "f(C)"
      ],
      "metadata": {
        "id": "SWuajOeJPs3X",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d5fb06d3-2016-4f57-b07d-e6e4321b6914"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([12, 15, 18], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using collectives"
      ],
      "metadata": {
        "id": "MCFHn3mauN1M"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(8)).reshape(2,4)\n",
        "arr"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-JgYIejauThB",
        "outputId": "99f76890-f026-4353-bbdf-8a3dca9b5540"
      },
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0, 1, 2, 3],\n",
              "       [4, 5, 6, 7]], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_pmap = jax.pmap(\n",
        "    jax.pmap(\n",
        "        lambda x: x/jax.lax.psum(x, axis_name=('rows','cols')),\n",
        "        axis_name='cols'\n",
        "    ),\n",
        "    axis_name='rows')"
      ],
      "metadata": {
        "id": "yU5SvNRYuPc7"
      },
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(n_pmap(arr))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HoHbmFEauVE8",
        "outputId": "3548968f-6a4d-41ac-d8d5-42bebfecd8b2"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_pmap(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DToC8YQqu8_w",
        "outputId": "34360bc9-009f-40e6-898b-e842d3984e29"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0.        , 0.03571429, 0.07142857, 0.10714287],\n",
              "       [0.14285715, 0.17857143, 0.21428573, 0.25      ]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_xmap = xmap(\n",
        "    lambda x: x/jax.lax.psum(x, axis_name=('rows','cols')),\n",
        "    in_axes=['rows', 'cols', ...],\n",
        "    out_axes=['rows', 'cols', ...]\n",
        ")"
      ],
      "metadata": {
        "id": "8QZoq0MruYL6"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.sum(n_xmap(arr))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SmsS7FRhuv1Q",
        "outputId": "1c9040a4-a012-46a1-ceb1-218305d5d422"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(1., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_xmap(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8QpuEN6luxmb",
        "outputId": "cb3574f0-e249-44c2-e9f0-398d7447aa9c"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0.        , 0.03571429, 0.07142857, 0.10714287],\n",
              "       [0.14285715, 0.17857143, 0.21428573, 0.25      ]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "arr = jnp.array(range(10000)).reshape(100,100)"
      ],
      "metadata": {
        "id": "xlAgK6buu7DI"
      },
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n_pmap(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 488
        },
        "id": "fcV0FTCEu1Cu",
        "outputId": "869afa75-1899-4fd4-dd44-0b68e7291d80"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[45], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mn_pmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43marr\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py:993\u001b[0m, in \u001b[0;36mUnloadedPmapExecutable.from_hlo\u001b[0;34m(hlo, pci, replicas, shards, tuple_args, unordered_effects, ordered_effects, host_callbacks, keepalive, jaxpr_debug_info, compiler_options)\u001b[0m\n\u001b[1;32m    990\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m shards\u001b[38;5;241m.\u001b[39mnum_global_shards \u001b[38;5;241m>\u001b[39m xb\u001b[38;5;241m.\u001b[39mdevice_count(pci\u001b[38;5;241m.\u001b[39mbackend):\n\u001b[1;32m    991\u001b[0m   msg \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompiling computation that requires \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m logical devices, but only \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m XLA \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    992\u001b[0m          \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdevices are available (num_replicas=\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 993\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(msg\u001b[38;5;241m.\u001b[39mformat(shards\u001b[38;5;241m.\u001b[39mnum_global_shards,\n\u001b[1;32m    994\u001b[0m                               xb\u001b[38;5;241m.\u001b[39mdevice_count(pci\u001b[38;5;241m.\u001b[39mbackend),\n\u001b[1;32m    995\u001b[0m                               replicas\u001b[38;5;241m.\u001b[39mnum_global_replicas))\n\u001b[1;32m    996\u001b[0m \u001b[38;5;66;03m# On a single host, we simply grab the first N devices from jax.devices().\u001b[39;00m\n\u001b[1;32m    997\u001b[0m \u001b[38;5;66;03m# In the single host case, we want the default device order of pmap to\u001b[39;00m\n\u001b[1;32m    998\u001b[0m \u001b[38;5;66;03m# match jax.devices().\u001b[39;00m\n\u001b[1;32m    999\u001b[0m \u001b[38;5;66;03m# On multiple hosts, we create a default device assignment that ensures\u001b[39;00m\n\u001b[1;32m   1000\u001b[0m \u001b[38;5;66;03m# each host is responsible for a contiguous set of replicas.\u001b[39;00m\n\u001b[1;32m   1001\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m shards\u001b[38;5;241m.\u001b[39mnum_global_shards \u001b[38;5;241m>\u001b[39m shards\u001b[38;5;241m.\u001b[39mnum_local_shards:\n\u001b[1;32m   1002\u001b[0m   \u001b[38;5;66;03m# TODO(skye): use a locality-aware assignment that satisfies the above\u001b[39;00m\n\u001b[1;32m   1003\u001b[0m   \u001b[38;5;66;03m# constraint.\u001b[39;00m\n",
            "\u001b[0;31mValueError\u001b[0m: compiling computation that requires 10000 logical devices, but only 8 XLA devices are available (num_replicas=10000)"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_xmap(arr)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xdmO8OpJvGR5",
        "outputId": "bc434b71-0d6d-4e8e-c8ec-2b8168b0132a"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[0.0000000e+00, 2.0002000e-08, 4.0004000e-08, ..., 1.9401939e-06,\n",
              "        1.9601960e-06, 1.9801980e-06],\n",
              "       [2.0002001e-06, 2.0202019e-06, 2.0402040e-06, ..., 3.9403940e-06,\n",
              "        3.9603960e-06, 3.9803981e-06],\n",
              "       [4.0004002e-06, 4.0204022e-06, 4.0404038e-06, ..., 5.9405938e-06,\n",
              "        5.9605959e-06, 5.9805980e-06],\n",
              "       ...,\n",
              "       [1.9401941e-04, 1.9403940e-04, 1.9405941e-04, ..., 1.9595960e-04,\n",
              "        1.9597959e-04, 1.9599960e-04],\n",
              "       [1.9601960e-04, 1.9603960e-04, 1.9605960e-04, ..., 1.9795979e-04,\n",
              "        1.9797980e-04, 1.9799981e-04],\n",
              "       [1.9801980e-04, 1.9803981e-04, 1.9805980e-04, ..., 1.9995999e-04,\n",
              "        1.9998000e-04, 1.9999999e-04]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 46
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using meshes"
      ],
      "metadata": {
        "id": "D_GmkPedvILO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from jax.sharding import Mesh"
      ],
      "metadata": {
        "id": "poE9ZsTevKdY"
      },
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np"
      ],
      "metadata": {
        "id": "SI0alZjP56SZ"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "jnp.array doesn't work for this type:"
      ],
      "metadata": {
        "id": "5-3l91_R8AK4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.array(jax.devices()).reshape(4, 2)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "fhVXuPX753TZ",
        "outputId": "21c0c0c7-c07c-4c58-f97b-5423a10c068a"
      },
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "error",
          "ename": "TypeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:528\u001b[0m, in \u001b[0;36mdtype\u001b[0;34m(x, canonicalize)\u001b[0m\n\u001b[1;32m    527\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 528\u001b[0m   dt \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult_type\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    529\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
            "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36mresult_type\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
            "\u001b[0;31mTypeError\u001b[0m: Cannot interpret 'TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)' as a data type",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:2015\u001b[0m, in \u001b[0;36marray\u001b[0;34m(object, dtype, copy, order, ndmin)\u001b[0m\n\u001b[1;32m   2014\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2015\u001b[0m   dtype \u001b[38;5;241m=\u001b[39m \u001b[43mdtypes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_lattice_result_type\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mleaves\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m leaves \u001b[38;5;28;01melse\u001b[39;00m dtypes\u001b[38;5;241m.\u001b[39mfloat_\n\u001b[1;32m   2016\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[1;32m   2017\u001b[0m   \u001b[38;5;66;03m# This happens if, e.g. one of the entries is a memoryview object.\u001b[39;00m\n\u001b[1;32m   2018\u001b[0m   \u001b[38;5;66;03m# This is rare, so we only handle it if the normal path fails.\u001b[39;00m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:537\u001b[0m, in \u001b[0;36m_lattice_result_type\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    536\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lattice_result_type\u001b[39m(\u001b[38;5;241m*\u001b[39margs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[DType, \u001b[38;5;28mbool\u001b[39m]:\n\u001b[0;32m--> 537\u001b[0m   dtypes, weak_types \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mzip\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m_dtype_and_weaktype\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43marg\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    538\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(dtypes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:537\u001b[0m, in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    536\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lattice_result_type\u001b[39m(\u001b[38;5;241m*\u001b[39margs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[DType, \u001b[38;5;28mbool\u001b[39m]:\n\u001b[0;32m--> 537\u001b[0m   dtypes, weak_types \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39m(\u001b[43m_dtype_and_weaktype\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args))\n\u001b[1;32m    538\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(dtypes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:360\u001b[0m, in \u001b[0;36m_dtype_and_weaktype\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    359\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a (dtype, weak_type) tuple for the given input.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 360\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdtype\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;28many\u001b[39m(value \u001b[38;5;129;01mis\u001b[39;00m typ \u001b[38;5;28;01mfor\u001b[39;00m typ \u001b[38;5;129;01min\u001b[39;00m _weak_types) \u001b[38;5;129;01mor\u001b[39;00m is_weakly_typed(value)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:530\u001b[0m, in \u001b[0;36mdtype\u001b[0;34m(x, canonicalize)\u001b[0m\n\u001b[1;32m    529\u001b[0m   \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 530\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot determine dtype of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[1;32m    531\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m dt \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m _jax_dtype_set \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_opaque_dtype(dt):\n",
            "\u001b[0;31mTypeError\u001b[0m: Cannot determine dtype of TPU_0(process=0,(0,0,0,0))",
            "\nDuring handling of the above exception, another exception occurred:\n",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[49], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mjax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevices\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m4\u001b[39m, \u001b[38;5;241m2\u001b[39m)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:2020\u001b[0m, in \u001b[0;36marray\u001b[0;34m(object, dtype, copy, order, ndmin)\u001b[0m\n\u001b[1;32m   2016\u001b[0m   \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[1;32m   2017\u001b[0m     \u001b[38;5;66;03m# This happens if, e.g. one of the entries is a memoryview object.\u001b[39;00m\n\u001b[1;32m   2018\u001b[0m     \u001b[38;5;66;03m# This is rare, so we only handle it if the normal path fails.\u001b[39;00m\n\u001b[1;32m   2019\u001b[0m     leaves \u001b[38;5;241m=\u001b[39m [_convert_to_array_if_dtype_fails(leaf) \u001b[38;5;28;01mfor\u001b[39;00m leaf \u001b[38;5;129;01min\u001b[39;00m leaves]\n\u001b[0;32m-> 2020\u001b[0m     dtype \u001b[38;5;241m=\u001b[39m \u001b[43mdtypes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_lattice_result_type\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mleaves\u001b[49m\u001b[43m)\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m   2022\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m weak_type:\n\u001b[1;32m   2023\u001b[0m   dtype \u001b[38;5;241m=\u001b[39m dtypes\u001b[38;5;241m.\u001b[39mcanonicalize_dtype(dtype, allow_opaque_dtype\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:537\u001b[0m, in \u001b[0;36m_lattice_result_type\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    536\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lattice_result_type\u001b[39m(\u001b[38;5;241m*\u001b[39margs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[DType, \u001b[38;5;28mbool\u001b[39m]:\n\u001b[0;32m--> 537\u001b[0m   dtypes, weak_types \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mzip\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m_dtype_and_weaktype\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43marg\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    538\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(dtypes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    539\u001b[0m     out_dtype \u001b[38;5;241m=\u001b[39m dtypes[\u001b[38;5;241m0\u001b[39m]\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:537\u001b[0m, in \u001b[0;36m<genexpr>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    536\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lattice_result_type\u001b[39m(\u001b[38;5;241m*\u001b[39margs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[DType, \u001b[38;5;28mbool\u001b[39m]:\n\u001b[0;32m--> 537\u001b[0m   dtypes, weak_types \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39m(\u001b[43m_dtype_and_weaktype\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args))\n\u001b[1;32m    538\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(dtypes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    539\u001b[0m     out_dtype \u001b[38;5;241m=\u001b[39m dtypes[\u001b[38;5;241m0\u001b[39m]\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:360\u001b[0m, in \u001b[0;36m_dtype_and_weaktype\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    358\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_dtype_and_weaktype\u001b[39m(value: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[DType, \u001b[38;5;28mbool\u001b[39m]:\n\u001b[1;32m    359\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124;03m\"\"\"Return a (dtype, weak_type) tuple for the given input.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 360\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdtype\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;28many\u001b[39m(value \u001b[38;5;129;01mis\u001b[39;00m typ \u001b[38;5;28;01mfor\u001b[39;00m typ \u001b[38;5;129;01min\u001b[39;00m _weak_types) \u001b[38;5;129;01mor\u001b[39;00m is_weakly_typed(value)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/dtypes.py:532\u001b[0m, in \u001b[0;36mdtype\u001b[0;34m(x, canonicalize)\u001b[0m\n\u001b[1;32m    530\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot determine dtype of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[1;32m    531\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m dt \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m _jax_dtype_set \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_opaque_dtype(dt):\n\u001b[0;32m--> 532\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mValue \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m with dtype \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mdt\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m is not a valid JAX array \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    533\u001b[0m                   \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtype. Only arrays of numeric types are supported by JAX.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    534\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m canonicalize_dtype(dt, allow_opaque_dtype\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mif\u001b[39;00m canonicalize \u001b[38;5;28;01melse\u001b[39;00m dt\n",
            "\u001b[0;31mTypeError\u001b[0m: Value 'TPU_0(process=0,(0,0,0,0))' with dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX."
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "devices = np.array(jax.devices()).reshape(4, 2)\n",
        "devices"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "26EiJI9V5vUA",
        "outputId": "24e09a55-387e-49c5-b0c6-c67785304a5f"
      },
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              "        TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1)],\n",
              "       [TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              "        TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1)],\n",
              "       [TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              "        TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1)],\n",
              "       [TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              "        TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]],\n",
              "      dtype=object)"
            ]
          },
          "metadata": {},
          "execution_count": 50
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "with Mesh(devices, ('x', 'y')):\n",
        "  n_xmap = xmap(\n",
        "    lambda x: x/jax.lax.psum(x, axis_name=('rows','cols')),\n",
        "    in_axes=['rows', 'cols', ...],\n",
        "    out_axes=['rows', 'cols', ...],\n",
        "    axis_resources={'rows': 'x', 'cols': 'y'}\n",
        "  )\n",
        "\n",
        "  res = n_xmap(arr)"
      ],
      "metadata": {
        "id": "Yy0yLNP96WVF"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "type(res), res.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PeX5vCz5DvLP",
        "outputId": "8ad80464-b5fa-474c-c17b-7230f3465137"
      },
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(jaxlib.xla_extension.ArrayImpl, (100, 100))"
            ]
          },
          "metadata": {},
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Simplifying initial xmap example (getting rid of reshaping)"
      ],
      "metadata": {
        "id": "ROMJFTHJ8Lu0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def dot(v1, v2):\n",
        "  return jnp.vdot(v1, v2)"
      ],
      "metadata": {
        "id": "nYWo9DGKEL74"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rng_key = random.PRNGKey(42)\n",
        "\n",
        "vs = random.normal(rng_key, shape=(20_000_000,3))\n",
        "v1s = vs[:10_000_000,:].T\n",
        "v2s = vs[10_000_000:,:].T\n",
        "\n",
        "v1s.shape, v2s.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hb0spwyQ8ZZT",
        "outputId": "57481595-5a8a-413d-c399-92645f60a635"
      },
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((3, 10000000), (3, 10000000))"
            ]
          },
          "metadata": {},
          "execution_count": 55
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "with Mesh(np.array(jax.devices()), ('device')):\n",
        "  f = xmap(dot,\n",
        "         in_axes=(\n",
        "             {1:'batch'},\n",
        "             {1:'batch'}\n",
        "         ),\n",
        "         out_axes=['batch', ...],\n",
        "         axis_resources={'batch': 'device'}\n",
        "  )\n",
        "  x_xmap=f(v1s,v2s)"
      ],
      "metadata": {
        "id": "gWsiDN7S8nOS"
      },
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x_xmap.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dunw8JQ79LVM",
        "outputId": "85ce042f-632e-4119-ad07-636a90b7972f"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(10000000,)"
            ]
          },
          "metadata": {},
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.numpy.all(x_xmap == x_pmap)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sC-QLQaL9v_Q",
        "outputId": "87d69f0d-d9a9-4bec-fdab-5338e5e895eb"
      },
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(True, dtype=bool)"
            ]
          },
          "metadata": {},
          "execution_count": 58
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Neural network example with xmap() [NOT WORKING]"
      ],
      "metadata": {
        "id": "DVXNIzeGNkaF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Preparing data"
      ],
      "metadata": {
        "id": "K0voYQTnSloK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mM6ZBCpqXvJ8",
        "outputId": "c15bfeea-7a3b-441d-9f0e-6f20df9b25e4"
      },
      "execution_count": 60,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\n",
            "Collecting tensorflow\n",
            "  Downloading tensorflow-2.13.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (479.6 MB)\n",
            "\u001b[K     |████████████████████            | 298.6 MB 129.0 MB/s eta 0:00:02"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "IOPub data rate exceeded.\n",
            "The Jupyter server will temporarily stop sending output\n",
            "to the client in order to avoid crashing it.\n",
            "To change this limit, set the config variable\n",
            "`--ServerApp.iopub_data_rate_limit`.\n",
            "\n",
            "Current values:\n",
            "ServerApp.iopub_data_rate_limit=1000000.0 (bytes/sec)\n",
            "ServerApp.rate_limit_window=3.0 (secs)\n",
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[K     |████████████████████████████████| 479.6 MB 9.5 kB/s \n",
            "\u001b[?25hCollecting absl-py>=1.0.0\n",
            "  Downloading absl_py-2.0.0-py3-none-any.whl (130 kB)\n",
            "\u001b[K     |████████████████████████████████| 130 kB 94.0 MB/s \n",
            "\u001b[?25hCollecting astunparse>=1.6.0\n",
            "  Downloading astunparse-1.6.3-py2.py3-none-any.whl (12 kB)\n",
            "Collecting flatbuffers>=23.1.21\n",
            "  Downloading flatbuffers-23.5.26-py2.py3-none-any.whl (26 kB)\n",
            "Collecting gast<=0.4.0,>=0.2.1\n",
            "  Downloading gast-0.4.0-py3-none-any.whl (9.8 kB)\n",
            "Collecting google-pasta>=0.1.1\n",
            "  Downloading google_pasta-0.2.0-py3-none-any.whl (57 kB)\n",
            "\u001b[K     |████████████████████████████████| 57 kB 8.0 MB/s \n",
            "\u001b[?25hCollecting grpcio<2.0,>=1.24.3\n",
            "  Downloading grpcio-1.59.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.3 MB 84.3 MB/s \n",
            "\u001b[?25hCollecting h5py>=2.9.0\n",
            "  Downloading h5py-3.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.8 MB 85.7 MB/s \n",
            "\u001b[?25hCollecting keras<2.14,>=2.13.1\n",
            "  Downloading keras-2.13.1-py3-none-any.whl (1.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.7 MB 89.7 MB/s \n",
            "\u001b[?25hCollecting libclang>=13.0.0\n",
            "  Downloading libclang-16.0.6-py2.py3-none-manylinux2010_x86_64.whl (22.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 22.9 MB 91.3 MB/s \n",
            "\u001b[?25hCollecting numpy<=1.24.3,>=1.22\n",
            "  Downloading numpy-1.24.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 17.3 MB 88.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: opt-einsum>=2.3.2 in ./.local/lib/python3.8/site-packages (from tensorflow) (3.3.0)\n",
            "Requirement already satisfied: packaging in /usr/lib/python3/dist-packages (from tensorflow) (20.3)\n",
            "Collecting protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3\n",
            "  Downloading protobuf-4.25.0-cp37-abi3-manylinux2014_x86_64.whl (294 kB)\n",
            "\u001b[K     |████████████████████████████████| 294 kB 100.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tensorflow) (62.3.2)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/lib/python3/dist-packages (from tensorflow) (1.14.0)\n",
            "Collecting tensorboard<2.14,>=2.13\n",
            "  Downloading tensorboard-2.13.0-py3-none-any.whl (5.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.6 MB 89.9 MB/s \n",
            "\u001b[?25hCollecting tensorflow-estimator<2.14,>=2.13.0\n",
            "  Downloading tensorflow_estimator-2.13.0-py2.py3-none-any.whl (440 kB)\n",
            "\u001b[K     |████████████████████████████████| 440 kB 106.4 MB/s \n",
            "\u001b[?25hCollecting termcolor>=1.1.0\n",
            "  Downloading termcolor-2.3.0-py3-none-any.whl (6.9 kB)\n",
            "Collecting typing-extensions<4.6.0,>=3.6.6\n",
            "  Downloading typing_extensions-4.5.0-py3-none-any.whl (27 kB)\n",
            "Collecting wrapt>=1.11.0\n",
            "  Downloading wrapt-1.15.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (81 kB)\n",
            "\u001b[K     |████████████████████████████████| 81 kB 12.4 MB/s \n",
            "\u001b[?25hCollecting tensorflow-io-gcs-filesystem>=0.23.1; platform_machine != \"arm64\" or platform_system != \"Darwin\"\n",
            "  Downloading tensorflow_io_gcs_filesystem-0.34.0-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (2.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.4 MB 87.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wheel<1.0,>=0.23.0 in /usr/lib/python3/dist-packages (from astunparse>=1.6.0->tensorflow) (0.34.2)\n",
            "Collecting google-auth<3,>=1.6.3\n",
            "  Downloading google_auth-2.23.4-py2.py3-none-any.whl (183 kB)\n",
            "\u001b[K     |████████████████████████████████| 183 kB 98.7 MB/s \n",
            "\u001b[?25hCollecting google-auth-oauthlib<1.1,>=0.5\n",
            "  Downloading google_auth_oauthlib-1.0.0-py2.py3-none-any.whl (18 kB)\n",
            "Collecting markdown>=2.6.8\n",
            "  Downloading Markdown-3.5.1-py3-none-any.whl (102 kB)\n",
            "\u001b[K     |████████████████████████████████| 102 kB 15.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests<3,>=2.21.0 in ./.local/lib/python3.8/site-packages (from tensorboard<2.14,>=2.13->tensorflow) (2.31.0)\n",
            "Collecting tensorboard-data-server<0.8.0,>=0.7.0\n",
            "  Downloading tensorboard_data_server-0.7.2-py3-none-any.whl (2.4 kB)\n",
            "Collecting werkzeug>=1.0.1\n",
            "  Downloading werkzeug-3.0.1-py3-none-any.whl (226 kB)\n",
            "\u001b[K     |████████████████████████████████| 226 kB 108.4 MB/s \n",
            "\u001b[?25hCollecting cachetools<6.0,>=2.0.0\n",
            "  Downloading cachetools-5.3.2-py3-none-any.whl (9.3 kB)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/lib/python3/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.14,>=2.13->tensorflow) (0.2.1)\n",
            "Collecting rsa<5,>=3.1.4\n",
            "  Downloading rsa-4.9-py3-none-any.whl (34 kB)\n",
            "Collecting requests-oauthlib>=0.7.0\n",
            "  Downloading requests_oauthlib-1.3.1-py2.py3-none-any.whl (23 kB)\n",
            "Requirement already satisfied: importlib-metadata>=4.4; python_version < \"3.10\" in ./.local/lib/python3.8/site-packages (from markdown>=2.6.8->tensorboard<2.14,>=2.13->tensorflow) (6.8.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (1.25.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests<3,>=2.21.0->tensorboard<2.14,>=2.13->tensorflow) (2019.11.28)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in ./.local/lib/python3.8/site-packages (from werkzeug>=1.0.1->tensorboard<2.14,>=2.13->tensorflow) (2.1.3)\n",
            "Requirement already satisfied: pyasn1>=0.1.3 in /usr/lib/python3/dist-packages (from rsa<5,>=3.1.4->google-auth<3,>=1.6.3->tensorboard<2.14,>=2.13->tensorflow) (0.4.2)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<1.1,>=0.5->tensorboard<2.14,>=2.13->tensorflow) (3.1.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/lib/python3/dist-packages (from importlib-metadata>=4.4; python_version < \"3.10\"->markdown>=2.6.8->tensorboard<2.14,>=2.13->tensorflow) (1.0.0)\n",
            "Installing collected packages: absl-py, astunparse, flatbuffers, gast, google-pasta, grpcio, numpy, h5py, keras, libclang, protobuf, cachetools, rsa, google-auth, requests-oauthlib, google-auth-oauthlib, markdown, tensorboard-data-server, werkzeug, tensorboard, tensorflow-estimator, termcolor, typing-extensions, wrapt, tensorflow-io-gcs-filesystem, tensorflow\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 1.24.4\n",
            "    Uninstalling numpy-1.24.4:\n",
            "      Successfully uninstalled numpy-1.24.4\n",
            "  Attempting uninstall: typing-extensions\n",
            "    Found existing installation: typing-extensions 4.8.0\n",
            "    Uninstalling typing-extensions-4.8.0:\n",
            "      Successfully uninstalled typing-extensions-4.8.0\n",
            "Successfully installed absl-py-2.0.0 astunparse-1.6.3 cachetools-5.3.2 flatbuffers-23.5.26 gast-0.4.0 google-auth-2.23.4 google-auth-oauthlib-1.0.0 google-pasta-0.2.0 grpcio-1.59.2 h5py-3.10.0 keras-2.13.1 libclang-16.0.6 markdown-3.5.1 numpy-1.24.3 protobuf-4.25.0 requests-oauthlib-1.3.1 rsa-4.9 tensorboard-2.13.0 tensorboard-data-server-0.7.2 tensorflow-2.13.1 tensorflow-estimator-2.13.0 tensorflow-io-gcs-filesystem-0.34.0 termcolor-2.3.0 typing-extensions-4.5.0 werkzeug-3.0.1 wrapt-1.15.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow_datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NSiCX_NcXxQW",
        "outputId": "746b63be-e9c0-4b90-8e2c-3826f1af9e2d"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.1.36ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\r\n",
            "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\r\n",
            "  warnings.warn(\n",
            "Collecting tensorflow_datasets\n",
            "  Downloading tensorflow_datasets-4.9.2-py3-none-any.whl (5.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.4 MB 4.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: absl-py in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.0.0)\n",
            "Collecting array-record\n",
            "  Downloading array_record-0.4.0-py38-none-any.whl (3.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.0 MB 92.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: click in /usr/lib/python3/dist-packages (from tensorflow_datasets) (7.0)\n",
            "Collecting dm-tree\n",
            "  Downloading dm_tree-0.1.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (152 kB)\n",
            "\u001b[K     |████████████████████████████████| 152 kB 101.1 MB/s \n",
            "\u001b[?25hCollecting etils[enp,epath]>=0.9.0\n",
            "  Downloading etils-1.3.0-py3-none-any.whl (126 kB)\n",
            "\u001b[K     |████████████████████████████████| 126 kB 109.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (1.24.3)\n",
            "Collecting promise\n",
            "  Downloading promise-2.3.tar.gz (19 kB)\n",
            "Requirement already satisfied: protobuf>=3.20 in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (4.25.0)\n",
            "Requirement already satisfied: psutil in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (5.9.6)\n",
            "Requirement already satisfied: requests>=2.19.0 in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.31.0)\n",
            "Collecting tensorflow-metadata\n",
            "  Downloading tensorflow_metadata-1.14.0-py3-none-any.whl (28 kB)\n",
            "Requirement already satisfied: termcolor in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (2.3.0)\n",
            "Collecting toml\n",
            "  Downloading toml-0.10.2-py2.py3-none-any.whl (16 kB)\n",
            "Collecting tqdm\n",
            "  Downloading tqdm-4.66.1-py3-none-any.whl (78 kB)\n",
            "\u001b[K     |████████████████████████████████| 78 kB 9.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wrapt in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (1.15.0)\n",
            "Requirement already satisfied: importlib-resources; python_version < \"3.9\" in ./.local/lib/python3.8/site-packages (from tensorflow_datasets) (6.1.0)\n",
            "Requirement already satisfied: typing_extensions; extra == \"epath\" in ./.local/lib/python3.8/site-packages (from etils[enp,epath]>=0.9.0->tensorflow_datasets) (4.5.0)\n",
            "Requirement already satisfied: zipp; extra == \"epath\" in /usr/lib/python3/dist-packages (from etils[enp,epath]>=0.9.0->tensorflow_datasets) (1.0.0)\n",
            "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from promise->tensorflow_datasets) (1.14.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.8/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (1.25.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests>=2.19.0->tensorflow_datasets) (2019.11.28)\n",
            "Collecting googleapis-common-protos<2,>=1.52.0\n",
            "  Downloading googleapis_common_protos-1.61.0-py2.py3-none-any.whl (230 kB)\n",
            "\u001b[K     |████████████████████████████████| 230 kB 102.4 MB/s \n",
            "\u001b[?25hBuilding wheels for collected packages: promise\n",
            "  Building wheel for promise (setup.py) ... \u001b[?25l-\b \b\\\b \b|\b \b/\b \b-\b \b\\\b \bdone\n",
            "\u001b[?25h  Created wheel for promise: filename=promise-2.3-py3-none-any.whl size=21485 sha256=5e79ed5a4461bb33ab12d34784f67b96b77e9fa0f642beb6e92d533727f1adeb\n",
            "  Stored in directory: /home/grigo/.cache/pip/wheels/54/aa/01/724885182f93150035a2a91bce34a12877e8067a97baaf5dc8\n",
            "Successfully built promise\n",
            "\u001b[31mERROR: tensorflow-metadata 1.14.0 has requirement absl-py<2.0.0,>=0.9, but you'll have absl-py 2.0.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow-metadata 1.14.0 has requirement protobuf<4.21,>=3.20.3, but you'll have protobuf 4.25.0 which is incompatible.\u001b[0m\n",
            "Installing collected packages: etils, array-record, dm-tree, promise, googleapis-common-protos, tensorflow-metadata, toml, tqdm, tensorflow-datasets\n",
            "Successfully installed array-record-0.4.0 dm-tree-0.1.8 etils-1.3.0 googleapis-common-protos-1.61.0 promise-2.3 tensorflow-datasets-4.9.2 tensorflow-metadata-1.14.0 toml-0.10.2 tqdm-4.66.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "data_dir = '/tmp/tfds'\n",
        "\n",
        "data, info = tfds.load(name=\"mnist\",\n",
        "                       data_dir=data_dir,\n",
        "                       as_supervised=True,\n",
        "                       with_info=True)\n",
        "\n",
        "data_train = data['train']\n",
        "data_test  = data['test']"
      ],
      "metadata": {
        "id": "M3B-FAZOOKOH",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0b87d625-dd90-44e4-c395-1bcdd8ff6df6"
      },
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2023-11-02 09:32:10.195256: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
            "/home/grigo/.local/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1mDownloading and preparing dataset 11.06 MiB (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /tmp/tfds/mnist/3.0.1...\u001b[0m\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Dl Completed...: 100%|████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 22.19 file/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1mDataset mnist downloaded and prepared to /tmp/tfds/mnist/3.0.1. Subsequent calls will reuse this data.\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "HEIGHT = 28\n",
        "WIDTH  = 28\n",
        "CHANNELS = 1\n",
        "NUM_PIXELS = HEIGHT * WIDTH * CHANNELS\n",
        "NUM_LABELS = info.features['label'].num_classes\n",
        "NUM_DEVICES = jax.device_count()\n",
        "BATCH_SIZE  = 32"
      ],
      "metadata": {
        "id": "GQW9KLknOSc4"
      },
      "execution_count": 63,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def preprocess(img, label):\n",
        "  \"\"\"Resize and preprocess images.\"\"\"\n",
        "  return (tf.cast(img, tf.float32)/255.0), label\n",
        "\n",
        "train_data = tfds.as_numpy(\n",
        "    data_train.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")\n",
        "test_data  = tfds.as_numpy(\n",
        "    data_test.map(preprocess).batch(NUM_DEVICES*BATCH_SIZE).prefetch(1)\n",
        ")"
      ],
      "metadata": {
        "id": "XEL-NZOjOV0_"
      },
      "execution_count": 64,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len(train_data)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bcwCdU2-GU80",
        "outputId": "fb6e68bf-a081-488b-aa1d-633369433faa"
      },
      "execution_count": 65,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "235"
            ]
          },
          "metadata": {},
          "execution_count": 65
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Preparing MLP"
      ],
      "metadata": {
        "id": "ilU3LzAeSolR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Potentially useful links:\n",
        "\n",
        "- my question https://github.com/google/jax/discussions/13861\n",
        "- translating simplified SPMD MLP to xmap (https://github.com/google/jax/issues/7167). Doesn't work because logsumexp uses pmax for which no differentiation rules implemented\n",
        "- some code for MLP with bias term and transformer blocks (https://gist.github.com/mattjj/ba9b24df446a90902d7b41aeb0766a99). Only xmap for loss, not xmap for diff(loss).\n",
        "- lax.pdot() documentation is actually missing (https://github.com/google/jax/pull/5020) (https://github.com/google/jax/discussions/13851)\n"
      ],
      "metadata": {
        "id": "45KlnzLtE1UB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jax import grad, jit, vmap, value_and_grad\n",
        "from jax import random\n",
        "from jax.nn import swish, logsumexp, one_hot"
      ],
      "metadata": {
        "id": "Q7B2PFjwOben"
      },
      "execution_count": 66,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "LAYER_SIZES = [28*28, 512, 10]\n",
        "AXES_NAMES  = ['inputs', 'hidden', 'classes']\n",
        "PARAM_SCALE = 0.01"
      ],
      "metadata": {
        "id": "aG1Ixe95OiIv"
      },
      "execution_count": 67,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def init_network_params(sizes, key=random.PRNGKey(0), scale=1e-2):\n",
        "  \"\"\"Initialize all layers for a fully-connected neural network with given sizes\"\"\"\n",
        "\n",
        "  def random_layer_params(m, n, key, scale=1e-2):\n",
        "    \"\"\"A helper function to randomly initialize weights and biases of a dense layer\"\"\"\n",
        "    w_key, b_key = random.split(key)\n",
        "    print(f'Generating layer params: w={(m,n)} b={(n,)}')\n",
        "    return scale * random.normal(w_key, (m, n)), scale * random.normal(b_key, (n,))\n",
        "\n",
        "  keys = random.split(key, len(sizes))\n",
        "  return [random_layer_params(m, n, k, scale) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n",
        "\n",
        "init_params = init_network_params(LAYER_SIZES, random.PRNGKey(0), scale=PARAM_SCALE)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8vENC9w64aR8",
        "outputId": "3aa09865-39e1-40dd-f31d-a62ca19f4258"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Generating layer params: w=(784, 512) b=(512,)\n",
            "Generating layer params: w=(512, 10) b=(10,)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def predict(params, image):\n",
        "  \"\"\"Function for per-example predictions.\"\"\"\n",
        "  activations = image\n",
        "  for (w,b), axis in zip(params[:-1], AXES_NAMES):\n",
        "    outputs = jax.lax.pdot(activations, w, axis) + b\n",
        "    activations = swish(outputs)\n",
        "\n",
        "  final_w, final_b = params[-1]\n",
        "  axis = AXES_NAMES[-2]\n",
        "  logits = jax.lax.pdot(activations, final_w, axis) + final_b\n",
        "  return logits"
      ],
      "metadata": {
        "id": "k9hQKmZXDC0p"
      },
      "execution_count": 69,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Loss and update functions"
      ],
      "metadata": {
        "id": "TdwS4P06Sr-P"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "INIT_LR = 1.0\n",
        "DECAY_RATE = 0.95\n",
        "DECAY_STEPS = 5\n",
        "NUM_EPOCHS  = 20"
      ],
      "metadata": {
        "id": "bwg667wGO3Yy"
      },
      "execution_count": 70,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def loss(params, images, labels):\n",
        "  \"\"\"Categorical cross entropy loss function.\"\"\"\n",
        "  logits = predict(params, images)\n",
        "  log_preds = logits - logsumexp(logits, AXES_NAMES[-1])\n",
        "  num_classes = jax.lax.psum(1, AXES_NAMES[-1])\n",
        "  targets = one_hot(labels, num_classes, axis=AXES_NAMES[-1])\n",
        "  losses = jax.lax.psum(targets*log_preds, AXES_NAMES[-1])\n",
        "  return -jax.lax.pmean(losses, \"batch\")"
      ],
      "metadata": {
        "id": "0Vmdg3ffadzW"
      },
      "execution_count": 71,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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"
      ],
      "metadata": {
        "id": "R8kvU4JrED54"
      },
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "update_named = xmap(update,\n",
        "                  in_axes=[\n",
        "                      [\n",
        "                          ({0: 'inputs', 1: 'hidden'}, {0: 'hidden'}),\n",
        "                          ({0: 'hidden', 1:'classes'}, {0:'classes'})\n",
        "                      ],\n",
        "                      {0: 'batch',  1: 'inputs'},\n",
        "                      {0: 'batch'},\n",
        "                      {}\n",
        "                  ],\n",
        "                  out_axes=(\n",
        "                      ([\n",
        "                        (['inputs', 'hidden', ...], ['hidden', ...]),\n",
        "                        (['hidden', 'classes', ...], ['classes', ...])\n",
        "                      ],\n",
        "                      {})\n",
        "                  ),\n",
        "                  )"
      ],
      "metadata": {
        "id": "GDM4aPOx7_BI"
      },
      "execution_count": 73,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Section for debugging purposes"
      ],
      "metadata": {
        "id": "GKMrN2vOPLiW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_data_iter = iter(train_data)\n",
        "x, y = next(train_data_iter)"
      ],
      "metadata": {
        "id": "9l5wT9U0O-7J"
      },
      "execution_count": 74,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x.shape, y.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fm2cdhekRNFC",
        "outputId": "a249288b-6130-43b8-da51-d4190c3d08d3"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((256, 28, 28, 1), (256,))"
            ]
          },
          "metadata": {},
          "execution_count": 75
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = jnp.reshape(x, (NUM_DEVICES*BATCH_SIZE, NUM_PIXELS))\n",
        "#y = jnp.reshape(one_hot(y, NUM_LABELS), (NUM_DEVICES*BATCH_SIZE, NUM_LABELS))\n",
        "x.shape, y.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ytR9XHd23525",
        "outputId": "89a1f5a3-c177-4681-8912-4f165bdc2441"
      },
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((256, 784), (256,))"
            ]
          },
          "metadata": {},
          "execution_count": 76
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "updated_params, loss_value = update_named(init_params, x, y, 0)"
      ],
      "metadata": {
        "id": "vSKemd9_QW1_",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "c25cef19-644a-420e-9774-64b18a24df7c"
      },
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NotImplementedError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m         Traceback (most recent call last)",
            "File \u001b[0;32m/usr/lib/python3.8/runpy.py:194\u001b[0m, in \u001b[0;36m_run_module_as_main\u001b[0;34m()\u001b[0m\n\u001b[1;32m    193\u001b[0m     sys\u001b[38;5;241m.\u001b[39margv[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m=\u001b[39m mod_spec\u001b[38;5;241m.\u001b[39morigin\n\u001b[0;32m--> 194\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _run_code(code, main_globals, \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    195\u001b[0m                  \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m, mod_spec)\n",
            "File \u001b[0;32m/usr/lib/python3.8/runpy.py:87\u001b[0m, in \u001b[0;36m_run_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m     80\u001b[0m run_globals\u001b[38;5;241m.\u001b[39mupdate(\u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m=\u001b[39m mod_name,\n\u001b[1;32m     81\u001b[0m                    \u001b[38;5;18m__file__\u001b[39m \u001b[38;5;241m=\u001b[39m fname,\n\u001b[1;32m     82\u001b[0m                    __cached__ \u001b[38;5;241m=\u001b[39m cached,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     85\u001b[0m                    __package__ \u001b[38;5;241m=\u001b[39m pkg_name,\n\u001b[1;32m     86\u001b[0m                    __spec__ \u001b[38;5;241m=\u001b[39m mod_spec)\n\u001b[0;32m---> 87\u001b[0m exec(code, run_globals)\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m run_globals\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel_launcher.py:17\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m app\u001b[38;5;241m.\u001b[39mlaunch_new_instance()\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/traitlets/config/application.py:1053\u001b[0m, in \u001b[0;36mlaunch_instance\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1052\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m-> 1053\u001b[0m app\u001b[38;5;241m.\u001b[39mstart()\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelapp.py:737\u001b[0m, in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 737\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mio_loop\u001b[38;5;241m.\u001b[39mstart()\n\u001b[1;32m    738\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/tornado/platform/asyncio.py:195\u001b[0m, in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    194\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstart\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 195\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop\u001b[38;5;241m.\u001b[39mrun_forever()\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mrun_forever\u001b[0;34m()\u001b[0m\n\u001b[1;32m    569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_once()\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36m_run_once\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m         handle\u001b[38;5;241m.\u001b[39m_run()\n\u001b[1;32m   1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/events.py:81\u001b[0m, in \u001b[0;36m_run\u001b[0;34m()\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_context\u001b[38;5;241m.\u001b[39mrun(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback, \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_args)\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:524\u001b[0m, in \u001b[0;36mdispatch_queue\u001b[0;34m()\u001b[0m\n\u001b[1;32m    523\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 524\u001b[0m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m    525\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:513\u001b[0m, in \u001b[0;36mprocess_one\u001b[0;34m()\u001b[0m\n\u001b[1;32m    512\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 513\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:418\u001b[0m, in \u001b[0;36mdispatch_shell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    417\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 418\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m    419\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:758\u001b[0m, in \u001b[0;36mexecute_request\u001b[0;34m()\u001b[0m\n\u001b[1;32m    757\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 758\u001b[0m     reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m    760\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:426\u001b[0m, in \u001b[0;36mdo_execute\u001b[0;34m()\u001b[0m\n\u001b[1;32m    425\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 426\u001b[0m     res \u001b[38;5;241m=\u001b[39m shell\u001b[38;5;241m.\u001b[39mrun_cell(\n\u001b[1;32m    427\u001b[0m         code,\n\u001b[1;32m    428\u001b[0m         store_history\u001b[38;5;241m=\u001b[39mstore_history,\n\u001b[1;32m    429\u001b[0m         silent\u001b[38;5;241m=\u001b[39msilent,\n\u001b[1;32m    430\u001b[0m         cell_id\u001b[38;5;241m=\u001b[39mcell_id,\n\u001b[1;32m    431\u001b[0m     )\n\u001b[1;32m    432\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/zmqshell.py:549\u001b[0m, in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    548\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 549\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39mrun_cell(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3009\u001b[0m, in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3008\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 3009\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_cell(\n\u001b[1;32m   3010\u001b[0m         raw_cell, store_history, silent, shell_futures, cell_id\n\u001b[1;32m   3011\u001b[0m     )\n\u001b[1;32m   3012\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3064\u001b[0m, in \u001b[0;36m_run_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3063\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 3064\u001b[0m     result \u001b[38;5;241m=\u001b[39m runner(coro)\n\u001b[1;32m   3065\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m()\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m     coro\u001b[38;5;241m.\u001b[39msend(\u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3269\u001b[0m, in \u001b[0;36mrun_cell_async\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3266\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3269\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m   3270\u001b[0m        interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m   3272\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3448\u001b[0m, in \u001b[0;36mrun_ast_nodes\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3447\u001b[0m     asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3448\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m   3449\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3508\u001b[0m, in \u001b[0;36mrun_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3507\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 3508\u001b[0m         exec(code_obj, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muser_global_ns, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muser_ns)\n\u001b[1;32m   3509\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   3510\u001b[0m     \u001b[38;5;66;03m# Reset our crash handler in place\u001b[39;00m\n",
            "Cell \u001b[0;32mIn[77], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m updated_params, loss_value \u001b[38;5;241m=\u001b[39m update_named(init_params, x, y, \u001b[38;5;241m0\u001b[39m)\n",
            "Cell \u001b[0;32mIn[72], line 2\u001b[0m, in \u001b[0;36mupdate\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdate\u001b[39m(params, x, y, epoch_number):\n\u001b[0;32m----> 2\u001b[0m   loss_value, grads \u001b[38;5;241m=\u001b[39m value_and_grad(loss)(params, x, y)\n\u001b[1;32m      3\u001b[0m   lr \u001b[38;5;241m=\u001b[39m INIT_LR \u001b[38;5;241m*\u001b[39m DECAY_RATE \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m (epoch_number \u001b[38;5;241m/\u001b[39m DECAY_STEPS)\n",
            "Cell \u001b[0;32mIn[71], line 4\u001b[0m, in \u001b[0;36mloss\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m logits \u001b[38;5;241m=\u001b[39m predict(params, images)\n\u001b[0;32m----> 4\u001b[0m log_preds \u001b[38;5;241m=\u001b[39m logits \u001b[38;5;241m-\u001b[39m logsumexp(logits, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m      5\u001b[0m num_classes \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mlax\u001b[38;5;241m.\u001b[39mpsum(\u001b[38;5;241m1\u001b[39m, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/ops/special.py:76\u001b[0m, in \u001b[0;36mlogsumexp\u001b[0;34m()\u001b[0m\n\u001b[1;32m     75\u001b[0m pos_dims, dims \u001b[38;5;241m=\u001b[39m _reduction_dims(a_arr, axis)\n\u001b[0;32m---> 76\u001b[0m amax \u001b[38;5;241m=\u001b[39m jnp\u001b[38;5;241m.\u001b[39mmax(a_arr, axis\u001b[38;5;241m=\u001b[39mdims, keepdims\u001b[38;5;241m=\u001b[39mkeepdims)\n\u001b[1;32m     77\u001b[0m amax \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39mstop_gradient(lax\u001b[38;5;241m.\u001b[39mselect(jnp\u001b[38;5;241m.\u001b[39misfinite(amax), amax, lax\u001b[38;5;241m.\u001b[39mfull_like(amax, \u001b[38;5;241m0\u001b[39m)))\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:256\u001b[0m, in \u001b[0;36mmax\u001b[0;34m()\u001b[0m\n\u001b[1;32m    252\u001b[0m \u001b[38;5;129m@_wraps\u001b[39m(np\u001b[38;5;241m.\u001b[39mmax, skip_params\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mout\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmax\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    254\u001b[0m         keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    255\u001b[0m         where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 256\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m _reduce_max(a, axis\u001b[38;5;241m=\u001b[39m_ensure_optional_axes(axis), out\u001b[38;5;241m=\u001b[39mout,\n\u001b[1;32m    257\u001b[0m                      keepdims\u001b[38;5;241m=\u001b[39mkeepdims, initial\u001b[38;5;241m=\u001b[39minitial, where\u001b[38;5;241m=\u001b[39mwhere)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:248\u001b[0m, in \u001b[0;36m_reduce_max\u001b[0;34m()\u001b[0m\n\u001b[1;32m    244\u001b[0m \u001b[38;5;129m@partial\u001b[39m(api\u001b[38;5;241m.\u001b[39mjit, static_argnames\u001b[38;5;241m=\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124maxis\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mkeepdims\u001b[39m\u001b[38;5;124m'\u001b[39m), inline\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    245\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_reduce_max\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    246\u001b[0m                 keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    247\u001b[0m                 where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 248\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m _reduction(a, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmax\u001b[39m\u001b[38;5;124m\"\u001b[39m, np\u001b[38;5;241m.\u001b[39mmax, lax\u001b[38;5;241m.\u001b[39mmax, \u001b[38;5;241m-\u001b[39mnp\u001b[38;5;241m.\u001b[39minf, has_identity\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    249\u001b[0m                     axis\u001b[38;5;241m=\u001b[39maxis, out\u001b[38;5;241m=\u001b[39mout, keepdims\u001b[38;5;241m=\u001b[39mkeepdims,\n\u001b[1;32m    250\u001b[0m                     initial\u001b[38;5;241m=\u001b[39minitial, where_\u001b[38;5;241m=\u001b[39mwhere, parallel_reduce\u001b[38;5;241m=\u001b[39mlax\u001b[38;5;241m.\u001b[39mpmax)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:130\u001b[0m, in \u001b[0;36m_reduction\u001b[0;34m()\u001b[0m\n\u001b[1;32m    129\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNamed reductions not implemented for jnp.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m()\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 130\u001b[0m   result \u001b[38;5;241m=\u001b[39m parallel_reduce(a, dims)\n\u001b[1;32m    131\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
            "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m: NotImplementedError: Differentiation rule for 'pmax' not implemented\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[77], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m updated_params, loss_value \u001b[38;5;241m=\u001b[39m \u001b[43mupdate_named\u001b[49m\u001b[43m(\u001b[49m\u001b[43minit_params\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:601\u001b[0m, in \u001b[0;36mxmap.<locals>.fun_mapped\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    599\u001b[0m tree_map(dispatch\u001b[38;5;241m.\u001b[39mcheck_arg, args)\n\u001b[1;32m    600\u001b[0m fun_flat, args_flat, params, _, out_tree \u001b[38;5;241m=\u001b[39m infer_params(\u001b[38;5;241m*\u001b[39margs)\n\u001b[0;32m--> 601\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxmap_p\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun_flat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    602\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m verify_outputs(out_flat, out_tree, params)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:823\u001b[0m, in \u001b[0;36mXMapPrimitive.bind\u001b[0;34m(self, fun, in_axes, *args, **params)\u001b[0m\n\u001b[1;32m    821\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, in_axes, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m    822\u001b[0m   \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(in_axes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(args), (in_axes, args)\n\u001b[0;32m--> 823\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcore\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43min_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/core.py:2393\u001b[0m, in \u001b[0;36mmap_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   2389\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmap_bind\u001b[39m(primitive: MapPrimitive, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m   2390\u001b[0m   map_bind_continuation, top_trace, fun, tracers, params \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m   2391\u001b[0m       map_bind_with_continuation(primitive, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams))\n\u001b[1;32m   2392\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m map_bind_continuation(\n\u001b[0;32m-> 2393\u001b[0m       \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:826\u001b[0m, in \u001b[0;36mXMapPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m    825\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m--> 826\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_xmap\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/core.py:818\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 818\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:630\u001b[0m, in \u001b[0;36mxmap_impl\u001b[0;34m(fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes, axis_resources, resource_env, backend, spmd_in_axes, spmd_out_axes_thunk, *args)\u001b[0m\n\u001b[1;32m    626\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mxmap_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, name, in_axes, out_axes_thunk, donated_invars,\n\u001b[1;32m    627\u001b[0m               global_axis_sizes, axis_resources, resource_env, backend,\n\u001b[1;32m    628\u001b[0m               spmd_in_axes, spmd_out_axes_thunk):\n\u001b[1;32m    629\u001b[0m   in_avals \u001b[38;5;241m=\u001b[39m [core\u001b[38;5;241m.\u001b[39mraise_to_shaped(core\u001b[38;5;241m.\u001b[39mget_aval(arg)) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[0;32m--> 630\u001b[0m   xmap_callable \u001b[38;5;241m=\u001b[39m \u001b[43mmake_xmap_callable\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    631\u001b[0m \u001b[43m      \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axes_thunk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mglobal_axis_sizes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    632\u001b[0m \u001b[43m      \u001b[49m\u001b[43maxis_resources\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresource_env\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    633\u001b[0m \u001b[43m      \u001b[49m\u001b[43mspmd_in_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mspmd_out_axes_thunk\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    634\u001b[0m \u001b[43m      \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mcompile()\u001b[38;5;241m.\u001b[39munsafe_call\n\u001b[1;32m    635\u001b[0m   distributed_debug_log((\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRunning xmapped function\u001b[39m\u001b[38;5;124m\"\u001b[39m, name),\n\u001b[1;32m    636\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpython function\u001b[39m\u001b[38;5;124m\"\u001b[39m, fun\u001b[38;5;241m.\u001b[39mf),\n\u001b[1;32m    637\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmesh\u001b[39m\u001b[38;5;124m\"\u001b[39m, resource_env\u001b[38;5;241m.\u001b[39mphysical_mesh),\n\u001b[1;32m    638\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mabstract args\u001b[39m\u001b[38;5;124m\"\u001b[39m, in_avals))\n\u001b[1;32m    639\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m xmap_callable(\u001b[38;5;241m*\u001b[39margs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/linear_util.py:345\u001b[0m, in \u001b[0;36mcache.<locals>.memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    343\u001b[0m   fun\u001b[38;5;241m.\u001b[39mpopulate_stores(stores)\n\u001b[1;32m    344\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 345\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    346\u001b[0m   cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n\u001b[1;32m    348\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ans\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:660\u001b[0m, in \u001b[0;36mmake_xmap_callable\u001b[0;34m(fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes, axis_resources, resource_env, backend, spmd_in_axes, spmd_out_axes_thunk, lowering_platform, *in_avals)\u001b[0m\n\u001b[1;32m    656\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m core\u001b[38;5;241m.\u001b[39mextend_axis_env_nd(global_axis_sizes\u001b[38;5;241m.\u001b[39mitems()):\n\u001b[1;32m    657\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m dispatch\u001b[38;5;241m.\u001b[39mlog_elapsed_time(\n\u001b[1;32m    658\u001b[0m       \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFinished tracing + transforming \u001b[39m\u001b[38;5;132;01m{fun_name}\u001b[39;00m\u001b[38;5;124m for xmap in \u001b[39m\u001b[38;5;132;01m{elapsed_time}\u001b[39;00m\u001b[38;5;124m sec\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    659\u001b[0m       fun_name\u001b[38;5;241m=\u001b[39mfun\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, event\u001b[38;5;241m=\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mJAXPR_TRACE_EVENT):\n\u001b[0;32m--> 660\u001b[0m     jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmapped_in_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    661\u001b[0m out_axes \u001b[38;5;241m=\u001b[39m out_axes_thunk()\n\u001b[1;32m    662\u001b[0m _check_out_avals_vs_out_axes(out_avals, out_axes, global_axis_sizes)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/profiler.py:314\u001b[0m, in \u001b[0;36mannotate_function.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    313\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m TraceAnnotation(name, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mdecorator_kwargs):\n\u001b[0;32m--> 314\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    315\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m wrapper\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py:2233\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, debug_info, keep_inputs)\u001b[0m\n\u001b[1;32m   2231\u001b[0m   main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m ()  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m   2232\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m core\u001b[38;5;241m.\u001b[39mnew_sublevel():\n\u001b[0;32m-> 2233\u001b[0m     jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2234\u001b[0m \u001b[43m      \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeep_inputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_inputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdebug_info\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdebug_info\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2235\u001b[0m   \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n\u001b[1;32m   2236\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m jaxpr, out_avals, consts\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py:2177\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals, keep_inputs, debug_info)\u001b[0m\n\u001b[1;32m   2175\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m _input_type_to_tracers(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[1;32m   2176\u001b[0m in_tracers_ \u001b[38;5;241m=\u001b[39m [t \u001b[38;5;28;01mfor\u001b[39;00m t, keep \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(in_tracers, keep_inputs) \u001b[38;5;28;01mif\u001b[39;00m keep]\n\u001b[0;32m-> 2177\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers_\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2178\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n\u001b[1;32m   2179\u001b[0m jaxpr, consts \u001b[38;5;241m=\u001b[39m frame\u001b[38;5;241m.\u001b[39mto_jaxpr(out_tracers)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/linear_util.py:188\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    185\u001b[0m gen \u001b[38;5;241m=\u001b[39m gen_static_args \u001b[38;5;241m=\u001b[39m out_store \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    187\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 188\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    190\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    191\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    192\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks\u001b[39;00m\n\u001b[1;32m    193\u001b[0m   \u001b[38;5;66;03m# only after this exception is handled, which can corrupt the global\u001b[39;00m\n\u001b[1;32m    194\u001b[0m   \u001b[38;5;66;03m# state.\u001b[39;00m\n\u001b[1;32m    195\u001b[0m   \u001b[38;5;28;01mwhile\u001b[39;00m stack:\n",
            "Cell \u001b[0;32mIn[72], line 2\u001b[0m, in \u001b[0;36mupdate\u001b[0;34m(params, x, y, epoch_number)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdate\u001b[39m(params, x, y, epoch_number):\n\u001b[0;32m----> 2\u001b[0m   loss_value, grads \u001b[38;5;241m=\u001b[39m \u001b[43mvalue_and_grad\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m   lr \u001b[38;5;241m=\u001b[39m INIT_LR \u001b[38;5;241m*\u001b[39m DECAY_RATE \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m (epoch_number \u001b[38;5;241m/\u001b[39m DECAY_STEPS)\n\u001b[1;32m      4\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m [(w \u001b[38;5;241m-\u001b[39m lr \u001b[38;5;241m*\u001b[39m dw, b \u001b[38;5;241m-\u001b[39m lr \u001b[38;5;241m*\u001b[39m db)\n\u001b[1;32m      5\u001b[0m           \u001b[38;5;28;01mfor\u001b[39;00m (w, b), (dw, db) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(params, grads)], loss_value\n",
            "    \u001b[0;31m[... skipping hidden 8 frame]\u001b[0m\n",
            "Cell \u001b[0;32mIn[71], line 4\u001b[0m, in \u001b[0;36mloss\u001b[0;34m(params, images, labels)\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Categorical cross entropy loss function.\"\"\"\u001b[39;00m\n\u001b[1;32m      3\u001b[0m logits \u001b[38;5;241m=\u001b[39m predict(params, images)\n\u001b[0;32m----> 4\u001b[0m log_preds \u001b[38;5;241m=\u001b[39m logits \u001b[38;5;241m-\u001b[39m \u001b[43mlogsumexp\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlogits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mAXES_NAMES\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m num_classes \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mlax\u001b[38;5;241m.\u001b[39mpsum(\u001b[38;5;241m1\u001b[39m, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m      6\u001b[0m targets \u001b[38;5;241m=\u001b[39m one_hot(labels, num_classes, axis\u001b[38;5;241m=\u001b[39mAXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/ops/special.py:76\u001b[0m, in \u001b[0;36mlogsumexp\u001b[0;34m(a, axis, b, keepdims, return_sign)\u001b[0m\n\u001b[1;32m     74\u001b[0m   b_arr \u001b[38;5;241m=\u001b[39m a_arr  \u001b[38;5;66;03m# for type checking\u001b[39;00m\n\u001b[1;32m     75\u001b[0m pos_dims, dims \u001b[38;5;241m=\u001b[39m _reduction_dims(a_arr, axis)\n\u001b[0;32m---> 76\u001b[0m amax \u001b[38;5;241m=\u001b[39m \u001b[43mjnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma_arr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeepdims\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeepdims\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     77\u001b[0m amax \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39mstop_gradient(lax\u001b[38;5;241m.\u001b[39mselect(jnp\u001b[38;5;241m.\u001b[39misfinite(amax), amax, lax\u001b[38;5;241m.\u001b[39mfull_like(amax, \u001b[38;5;241m0\u001b[39m)))\n\u001b[1;32m     78\u001b[0m amax_with_dims \u001b[38;5;241m=\u001b[39m amax \u001b[38;5;28;01mif\u001b[39;00m keepdims \u001b[38;5;28;01melse\u001b[39;00m lax\u001b[38;5;241m.\u001b[39mexpand_dims(amax, pos_dims)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:256\u001b[0m, in \u001b[0;36mmax\u001b[0;34m(a, axis, out, keepdims, initial, where)\u001b[0m\n\u001b[1;32m    252\u001b[0m \u001b[38;5;129m@_wraps\u001b[39m(np\u001b[38;5;241m.\u001b[39mmax, skip_params\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mout\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmax\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    254\u001b[0m         keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    255\u001b[0m         where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 256\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_reduce_max\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_ensure_optional_axes\u001b[49m\u001b[43m(\u001b[49m\u001b[43maxis\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    257\u001b[0m \u001b[43m                     \u001b[49m\u001b[43mkeepdims\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeepdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minitial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minitial\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwhere\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwhere\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 17 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/ad.py:314\u001b[0m, in \u001b[0;36mJVPTrace.process_primitive\u001b[0;34m(self, primitive, tracers, params)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m jvp:\n\u001b[1;32m    313\u001b[0m   msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDifferentiation rule for \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprimitive\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m not implemented\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m--> 314\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(msg)\n\u001b[1;32m    315\u001b[0m primal_out, tangent_out \u001b[38;5;241m=\u001b[39m jvp(primals_in, tangents_in, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams)\n\u001b[1;32m    316\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m primitive\u001b[38;5;241m.\u001b[39mmultiple_results:\n",
            "\u001b[0;31mNotImplementedError\u001b[0m: Differentiation rule for 'pmax' not implemented"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "loss_value"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 175
        },
        "id": "QqMETg4EBdNC",
        "outputId": "a7308812-30a7-466f-c0e4-cb0c2ba37c39"
      },
      "execution_count": 78,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[78], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mloss_value\u001b[49m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'loss_value' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "?jax.lax.pdot"
      ],
      "metadata": {
        "id": "aCDz6vpire4R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training loop"
      ],
      "metadata": {
        "id": "4PFlq7EcS3q_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@jit\n",
        "def batch_accuracy(params, images, targets):\n",
        "  images = jnp.reshape(images, (len(images), NUM_PIXELS))\n",
        "  predicted_class = jnp.argmax(vmap(predict)(params, images), axis=1)\n",
        "  return jnp.mean(predicted_class == targets)\n",
        "\n",
        "def accuracy(params, data):\n",
        "  accs = []\n",
        "  for images, targets in data:\n",
        "    accs.append(batch_accuracy(params, images, targets))\n",
        "  return jnp.mean(jnp.array(accs))"
      ],
      "metadata": {
        "id": "Zw4DIGl5InbS"
      },
      "execution_count": 79,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "params = init_params\n",
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "  start_time = time.time()\n",
        "  losses = []\n",
        "  for x, y in train_data:\n",
        "    num_elements = len(y)\n",
        "    x = jnp.reshape(x, (num_elements, NUM_PIXELS))\n",
        "    #y = jnp.reshape(one_hot(y, NUM_LABELS), (NUM_DEVICES, num_elements//NUM_DEVICES, NUM_LABELS))\n",
        "    params, loss_value = update_named(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))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "skYy98qTFUpP",
        "outputId": "52f292f8-fb9a-4c31-9c92-b003a297ce85"
      },
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2023-11-02 09:33:16.226920: W tensorflow/core/kernels/data/cache_dataset_ops.cc:854] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "NotImplementedError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m         Traceback (most recent call last)",
            "File \u001b[0;32m/usr/lib/python3.8/runpy.py:194\u001b[0m, in \u001b[0;36m_run_module_as_main\u001b[0;34m()\u001b[0m\n\u001b[1;32m    193\u001b[0m     sys\u001b[38;5;241m.\u001b[39margv[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m=\u001b[39m mod_spec\u001b[38;5;241m.\u001b[39morigin\n\u001b[0;32m--> 194\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _run_code(code, main_globals, \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    195\u001b[0m                  \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m, mod_spec)\n",
            "File \u001b[0;32m/usr/lib/python3.8/runpy.py:87\u001b[0m, in \u001b[0;36m_run_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m     80\u001b[0m run_globals\u001b[38;5;241m.\u001b[39mupdate(\u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m=\u001b[39m mod_name,\n\u001b[1;32m     81\u001b[0m                    \u001b[38;5;18m__file__\u001b[39m \u001b[38;5;241m=\u001b[39m fname,\n\u001b[1;32m     82\u001b[0m                    __cached__ \u001b[38;5;241m=\u001b[39m cached,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     85\u001b[0m                    __package__ \u001b[38;5;241m=\u001b[39m pkg_name,\n\u001b[1;32m     86\u001b[0m                    __spec__ \u001b[38;5;241m=\u001b[39m mod_spec)\n\u001b[0;32m---> 87\u001b[0m exec(code, run_globals)\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m run_globals\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel_launcher.py:17\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m app\u001b[38;5;241m.\u001b[39mlaunch_new_instance()\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/traitlets/config/application.py:1053\u001b[0m, in \u001b[0;36mlaunch_instance\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1052\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m-> 1053\u001b[0m app\u001b[38;5;241m.\u001b[39mstart()\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelapp.py:737\u001b[0m, in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    736\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 737\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mio_loop\u001b[38;5;241m.\u001b[39mstart()\n\u001b[1;32m    738\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/tornado/platform/asyncio.py:195\u001b[0m, in \u001b[0;36mstart\u001b[0;34m()\u001b[0m\n\u001b[1;32m    194\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mstart\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 195\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop\u001b[38;5;241m.\u001b[39mrun_forever()\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mrun_forever\u001b[0;34m()\u001b[0m\n\u001b[1;32m    569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_once()\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36m_run_once\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m         handle\u001b[38;5;241m.\u001b[39m_run()\n\u001b[1;32m   1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
            "File \u001b[0;32m/usr/lib/python3.8/asyncio/events.py:81\u001b[0m, in \u001b[0;36m_run\u001b[0;34m()\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_context\u001b[38;5;241m.\u001b[39mrun(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback, \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_args)\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:524\u001b[0m, in \u001b[0;36mdispatch_queue\u001b[0;34m()\u001b[0m\n\u001b[1;32m    523\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 524\u001b[0m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m    525\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:513\u001b[0m, in \u001b[0;36mprocess_one\u001b[0;34m()\u001b[0m\n\u001b[1;32m    512\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 513\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:418\u001b[0m, in \u001b[0;36mdispatch_shell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    417\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 418\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m    419\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/kernelbase.py:758\u001b[0m, in \u001b[0;36mexecute_request\u001b[0;34m()\u001b[0m\n\u001b[1;32m    757\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 758\u001b[0m     reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m    760\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/ipkernel.py:426\u001b[0m, in \u001b[0;36mdo_execute\u001b[0;34m()\u001b[0m\n\u001b[1;32m    425\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 426\u001b[0m     res \u001b[38;5;241m=\u001b[39m shell\u001b[38;5;241m.\u001b[39mrun_cell(\n\u001b[1;32m    427\u001b[0m         code,\n\u001b[1;32m    428\u001b[0m         store_history\u001b[38;5;241m=\u001b[39mstore_history,\n\u001b[1;32m    429\u001b[0m         silent\u001b[38;5;241m=\u001b[39msilent,\n\u001b[1;32m    430\u001b[0m         cell_id\u001b[38;5;241m=\u001b[39mcell_id,\n\u001b[1;32m    431\u001b[0m     )\n\u001b[1;32m    432\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/ipykernel/zmqshell.py:549\u001b[0m, in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m    548\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 549\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39mrun_cell(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3009\u001b[0m, in \u001b[0;36mrun_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3008\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 3009\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_run_cell(\n\u001b[1;32m   3010\u001b[0m         raw_cell, store_history, silent, shell_futures, cell_id\n\u001b[1;32m   3011\u001b[0m     )\n\u001b[1;32m   3012\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3064\u001b[0m, in \u001b[0;36m_run_cell\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3063\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 3064\u001b[0m     result \u001b[38;5;241m=\u001b[39m runner(coro)\n\u001b[1;32m   3065\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m()\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m     coro\u001b[38;5;241m.\u001b[39msend(\u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3269\u001b[0m, in \u001b[0;36mrun_cell_async\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3266\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3269\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m   3270\u001b[0m        interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m   3272\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3448\u001b[0m, in \u001b[0;36mrun_ast_nodes\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3447\u001b[0m     asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3448\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m   3449\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3508\u001b[0m, in \u001b[0;36mrun_code\u001b[0;34m()\u001b[0m\n\u001b[1;32m   3507\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 3508\u001b[0m         exec(code_obj, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muser_global_ns, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39muser_ns)\n\u001b[1;32m   3509\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m   3510\u001b[0m     \u001b[38;5;66;03m# Reset our crash handler in place\u001b[39;00m\n",
            "Cell \u001b[0;32mIn[77], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m updated_params, loss_value \u001b[38;5;241m=\u001b[39m update_named(init_params, x, y, \u001b[38;5;241m0\u001b[39m)\n",
            "Cell \u001b[0;32mIn[72], line 2\u001b[0m, in \u001b[0;36mupdate\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdate\u001b[39m(params, x, y, epoch_number):\n\u001b[0;32m----> 2\u001b[0m   loss_value, grads \u001b[38;5;241m=\u001b[39m value_and_grad(loss)(params, x, y)\n\u001b[1;32m      3\u001b[0m   lr \u001b[38;5;241m=\u001b[39m INIT_LR \u001b[38;5;241m*\u001b[39m DECAY_RATE \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m (epoch_number \u001b[38;5;241m/\u001b[39m DECAY_STEPS)\n",
            "Cell \u001b[0;32mIn[71], line 4\u001b[0m, in \u001b[0;36mloss\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m logits \u001b[38;5;241m=\u001b[39m predict(params, images)\n\u001b[0;32m----> 4\u001b[0m log_preds \u001b[38;5;241m=\u001b[39m logits \u001b[38;5;241m-\u001b[39m logsumexp(logits, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m      5\u001b[0m num_classes \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mlax\u001b[38;5;241m.\u001b[39mpsum(\u001b[38;5;241m1\u001b[39m, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/ops/special.py:76\u001b[0m, in \u001b[0;36mlogsumexp\u001b[0;34m()\u001b[0m\n\u001b[1;32m     75\u001b[0m pos_dims, dims \u001b[38;5;241m=\u001b[39m _reduction_dims(a_arr, axis)\n\u001b[0;32m---> 76\u001b[0m amax \u001b[38;5;241m=\u001b[39m jnp\u001b[38;5;241m.\u001b[39mmax(a_arr, axis\u001b[38;5;241m=\u001b[39mdims, keepdims\u001b[38;5;241m=\u001b[39mkeepdims)\n\u001b[1;32m     77\u001b[0m amax \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39mstop_gradient(lax\u001b[38;5;241m.\u001b[39mselect(jnp\u001b[38;5;241m.\u001b[39misfinite(amax), amax, lax\u001b[38;5;241m.\u001b[39mfull_like(amax, \u001b[38;5;241m0\u001b[39m)))\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:256\u001b[0m, in \u001b[0;36mmax\u001b[0;34m()\u001b[0m\n\u001b[1;32m    252\u001b[0m \u001b[38;5;129m@_wraps\u001b[39m(np\u001b[38;5;241m.\u001b[39mmax, skip_params\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mout\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmax\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    254\u001b[0m         keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    255\u001b[0m         where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 256\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m _reduce_max(a, axis\u001b[38;5;241m=\u001b[39m_ensure_optional_axes(axis), out\u001b[38;5;241m=\u001b[39mout,\n\u001b[1;32m    257\u001b[0m                      keepdims\u001b[38;5;241m=\u001b[39mkeepdims, initial\u001b[38;5;241m=\u001b[39minitial, where\u001b[38;5;241m=\u001b[39mwhere)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:248\u001b[0m, in \u001b[0;36m_reduce_max\u001b[0;34m()\u001b[0m\n\u001b[1;32m    244\u001b[0m \u001b[38;5;129m@partial\u001b[39m(api\u001b[38;5;241m.\u001b[39mjit, static_argnames\u001b[38;5;241m=\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124maxis\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mkeepdims\u001b[39m\u001b[38;5;124m'\u001b[39m), inline\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    245\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_reduce_max\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    246\u001b[0m                 keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    247\u001b[0m                 where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 248\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m _reduction(a, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmax\u001b[39m\u001b[38;5;124m\"\u001b[39m, np\u001b[38;5;241m.\u001b[39mmax, lax\u001b[38;5;241m.\u001b[39mmax, \u001b[38;5;241m-\u001b[39mnp\u001b[38;5;241m.\u001b[39minf, has_identity\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    249\u001b[0m                     axis\u001b[38;5;241m=\u001b[39maxis, out\u001b[38;5;241m=\u001b[39mout, keepdims\u001b[38;5;241m=\u001b[39mkeepdims,\n\u001b[1;32m    250\u001b[0m                     initial\u001b[38;5;241m=\u001b[39minitial, where_\u001b[38;5;241m=\u001b[39mwhere, parallel_reduce\u001b[38;5;241m=\u001b[39mlax\u001b[38;5;241m.\u001b[39mpmax)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:130\u001b[0m, in \u001b[0;36m_reduction\u001b[0;34m()\u001b[0m\n\u001b[1;32m    129\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mNamed reductions not implemented for jnp.\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m()\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 130\u001b[0m   result \u001b[38;5;241m=\u001b[39m parallel_reduce(a, dims)\n\u001b[1;32m    131\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
            "\u001b[0;31mJaxStackTraceBeforeTransformation\u001b[0m: NotImplementedError: Differentiation rule for 'pmax' not implemented\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[80], line 12\u001b[0m\n\u001b[1;32m     10\u001b[0m   x \u001b[38;5;241m=\u001b[39m jnp\u001b[38;5;241m.\u001b[39mreshape(x, (num_elements, NUM_PIXELS))\n\u001b[1;32m     11\u001b[0m   \u001b[38;5;66;03m#y = jnp.reshape(one_hot(y, NUM_LABELS), (NUM_DEVICES, num_elements//NUM_DEVICES, NUM_LABELS))\u001b[39;00m\n\u001b[0;32m---> 12\u001b[0m   params, loss_value \u001b[38;5;241m=\u001b[39m \u001b[43mupdate_named\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     13\u001b[0m   losses\u001b[38;5;241m.\u001b[39mappend(loss_value)\n\u001b[1;32m     14\u001b[0m epoch_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m-\u001b[39m start_time\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:601\u001b[0m, in \u001b[0;36mxmap.<locals>.fun_mapped\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    599\u001b[0m tree_map(dispatch\u001b[38;5;241m.\u001b[39mcheck_arg, args)\n\u001b[1;32m    600\u001b[0m fun_flat, args_flat, params, _, out_tree \u001b[38;5;241m=\u001b[39m infer_params(\u001b[38;5;241m*\u001b[39margs)\n\u001b[0;32m--> 601\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxmap_p\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun_flat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    602\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m verify_outputs(out_flat, out_tree, params)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:823\u001b[0m, in \u001b[0;36mXMapPrimitive.bind\u001b[0;34m(self, fun, in_axes, *args, **params)\u001b[0m\n\u001b[1;32m    821\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, in_axes, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m    822\u001b[0m   \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(in_axes) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(args), (in_axes, args)\n\u001b[0;32m--> 823\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcore\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmap_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43min_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/core.py:2393\u001b[0m, in \u001b[0;36mmap_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   2389\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmap_bind\u001b[39m(primitive: MapPrimitive, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m   2390\u001b[0m   map_bind_continuation, top_trace, fun, tracers, params \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m   2391\u001b[0m       map_bind_with_continuation(primitive, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams))\n\u001b[1;32m   2392\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m map_bind_continuation(\n\u001b[0;32m-> 2393\u001b[0m       \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:826\u001b[0m, in \u001b[0;36mXMapPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m    825\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m--> 826\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_xmap\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/core.py:818\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 818\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:630\u001b[0m, in \u001b[0;36mxmap_impl\u001b[0;34m(fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes, axis_resources, resource_env, backend, spmd_in_axes, spmd_out_axes_thunk, *args)\u001b[0m\n\u001b[1;32m    626\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mxmap_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, name, in_axes, out_axes_thunk, donated_invars,\n\u001b[1;32m    627\u001b[0m               global_axis_sizes, axis_resources, resource_env, backend,\n\u001b[1;32m    628\u001b[0m               spmd_in_axes, spmd_out_axes_thunk):\n\u001b[1;32m    629\u001b[0m   in_avals \u001b[38;5;241m=\u001b[39m [core\u001b[38;5;241m.\u001b[39mraise_to_shaped(core\u001b[38;5;241m.\u001b[39mget_aval(arg)) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[0;32m--> 630\u001b[0m   xmap_callable \u001b[38;5;241m=\u001b[39m \u001b[43mmake_xmap_callable\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    631\u001b[0m \u001b[43m      \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_axes_thunk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mglobal_axis_sizes\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    632\u001b[0m \u001b[43m      \u001b[49m\u001b[43maxis_resources\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresource_env\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    633\u001b[0m \u001b[43m      \u001b[49m\u001b[43mspmd_in_axes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mspmd_out_axes_thunk\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    634\u001b[0m \u001b[43m      \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mcompile()\u001b[38;5;241m.\u001b[39munsafe_call\n\u001b[1;32m    635\u001b[0m   distributed_debug_log((\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRunning xmapped function\u001b[39m\u001b[38;5;124m\"\u001b[39m, name),\n\u001b[1;32m    636\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpython function\u001b[39m\u001b[38;5;124m\"\u001b[39m, fun\u001b[38;5;241m.\u001b[39mf),\n\u001b[1;32m    637\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmesh\u001b[39m\u001b[38;5;124m\"\u001b[39m, resource_env\u001b[38;5;241m.\u001b[39mphysical_mesh),\n\u001b[1;32m    638\u001b[0m                         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mabstract args\u001b[39m\u001b[38;5;124m\"\u001b[39m, in_avals))\n\u001b[1;32m    639\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m xmap_callable(\u001b[38;5;241m*\u001b[39margs)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/linear_util.py:345\u001b[0m, in \u001b[0;36mcache.<locals>.memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    343\u001b[0m   fun\u001b[38;5;241m.\u001b[39mpopulate_stores(stores)\n\u001b[1;32m    344\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 345\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    346\u001b[0m   cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n\u001b[1;32m    348\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ans\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/maps.py:660\u001b[0m, in \u001b[0;36mmake_xmap_callable\u001b[0;34m(fun, name, in_axes, out_axes_thunk, donated_invars, global_axis_sizes, axis_resources, resource_env, backend, spmd_in_axes, spmd_out_axes_thunk, lowering_platform, *in_avals)\u001b[0m\n\u001b[1;32m    656\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m core\u001b[38;5;241m.\u001b[39mextend_axis_env_nd(global_axis_sizes\u001b[38;5;241m.\u001b[39mitems()):\n\u001b[1;32m    657\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m dispatch\u001b[38;5;241m.\u001b[39mlog_elapsed_time(\n\u001b[1;32m    658\u001b[0m       \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFinished tracing + transforming \u001b[39m\u001b[38;5;132;01m{fun_name}\u001b[39;00m\u001b[38;5;124m for xmap in \u001b[39m\u001b[38;5;132;01m{elapsed_time}\u001b[39;00m\u001b[38;5;124m sec\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    659\u001b[0m       fun_name\u001b[38;5;241m=\u001b[39mfun\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, event\u001b[38;5;241m=\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mJAXPR_TRACE_EVENT):\n\u001b[0;32m--> 660\u001b[0m     jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmapped_in_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    661\u001b[0m out_axes \u001b[38;5;241m=\u001b[39m out_axes_thunk()\n\u001b[1;32m    662\u001b[0m _check_out_avals_vs_out_axes(out_avals, out_axes, global_axis_sizes)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/profiler.py:314\u001b[0m, in \u001b[0;36mannotate_function.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    313\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m TraceAnnotation(name, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mdecorator_kwargs):\n\u001b[0;32m--> 314\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    315\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m wrapper\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py:2233\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, debug_info, keep_inputs)\u001b[0m\n\u001b[1;32m   2231\u001b[0m   main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m ()  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m   2232\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m core\u001b[38;5;241m.\u001b[39mnew_sublevel():\n\u001b[0;32m-> 2233\u001b[0m     jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2234\u001b[0m \u001b[43m      \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeep_inputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeep_inputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdebug_info\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdebug_info\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2235\u001b[0m   \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n\u001b[1;32m   2236\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m jaxpr, out_avals, consts\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py:2177\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals, keep_inputs, debug_info)\u001b[0m\n\u001b[1;32m   2175\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m _input_type_to_tracers(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[1;32m   2176\u001b[0m in_tracers_ \u001b[38;5;241m=\u001b[39m [t \u001b[38;5;28;01mfor\u001b[39;00m t, keep \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(in_tracers, keep_inputs) \u001b[38;5;28;01mif\u001b[39;00m keep]\n\u001b[0;32m-> 2177\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers_\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2178\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n\u001b[1;32m   2179\u001b[0m jaxpr, consts \u001b[38;5;241m=\u001b[39m frame\u001b[38;5;241m.\u001b[39mto_jaxpr(out_tracers)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/linear_util.py:188\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    185\u001b[0m gen \u001b[38;5;241m=\u001b[39m gen_static_args \u001b[38;5;241m=\u001b[39m out_store \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    187\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 188\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    190\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    191\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    192\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks\u001b[39;00m\n\u001b[1;32m    193\u001b[0m   \u001b[38;5;66;03m# only after this exception is handled, which can corrupt the global\u001b[39;00m\n\u001b[1;32m    194\u001b[0m   \u001b[38;5;66;03m# state.\u001b[39;00m\n\u001b[1;32m    195\u001b[0m   \u001b[38;5;28;01mwhile\u001b[39;00m stack:\n",
            "Cell \u001b[0;32mIn[72], line 2\u001b[0m, in \u001b[0;36mupdate\u001b[0;34m(params, x, y, epoch_number)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdate\u001b[39m(params, x, y, epoch_number):\n\u001b[0;32m----> 2\u001b[0m   loss_value, grads \u001b[38;5;241m=\u001b[39m \u001b[43mvalue_and_grad\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m   lr \u001b[38;5;241m=\u001b[39m INIT_LR \u001b[38;5;241m*\u001b[39m DECAY_RATE \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m (epoch_number \u001b[38;5;241m/\u001b[39m DECAY_STEPS)\n\u001b[1;32m      4\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m [(w \u001b[38;5;241m-\u001b[39m lr \u001b[38;5;241m*\u001b[39m dw, b \u001b[38;5;241m-\u001b[39m lr \u001b[38;5;241m*\u001b[39m db)\n\u001b[1;32m      5\u001b[0m           \u001b[38;5;28;01mfor\u001b[39;00m (w, b), (dw, db) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(params, grads)], loss_value\n",
            "    \u001b[0;31m[... skipping hidden 8 frame]\u001b[0m\n",
            "Cell \u001b[0;32mIn[71], line 4\u001b[0m, in \u001b[0;36mloss\u001b[0;34m(params, images, labels)\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Categorical cross entropy loss function.\"\"\"\u001b[39;00m\n\u001b[1;32m      3\u001b[0m logits \u001b[38;5;241m=\u001b[39m predict(params, images)\n\u001b[0;32m----> 4\u001b[0m log_preds \u001b[38;5;241m=\u001b[39m logits \u001b[38;5;241m-\u001b[39m \u001b[43mlogsumexp\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlogits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mAXES_NAMES\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      5\u001b[0m num_classes \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mlax\u001b[38;5;241m.\u001b[39mpsum(\u001b[38;5;241m1\u001b[39m, AXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m      6\u001b[0m targets \u001b[38;5;241m=\u001b[39m one_hot(labels, num_classes, axis\u001b[38;5;241m=\u001b[39mAXES_NAMES[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/ops/special.py:76\u001b[0m, in \u001b[0;36mlogsumexp\u001b[0;34m(a, axis, b, keepdims, return_sign)\u001b[0m\n\u001b[1;32m     74\u001b[0m   b_arr \u001b[38;5;241m=\u001b[39m a_arr  \u001b[38;5;66;03m# for type checking\u001b[39;00m\n\u001b[1;32m     75\u001b[0m pos_dims, dims \u001b[38;5;241m=\u001b[39m _reduction_dims(a_arr, axis)\n\u001b[0;32m---> 76\u001b[0m amax \u001b[38;5;241m=\u001b[39m \u001b[43mjnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma_arr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeepdims\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeepdims\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     77\u001b[0m amax \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39mstop_gradient(lax\u001b[38;5;241m.\u001b[39mselect(jnp\u001b[38;5;241m.\u001b[39misfinite(amax), amax, lax\u001b[38;5;241m.\u001b[39mfull_like(amax, \u001b[38;5;241m0\u001b[39m)))\n\u001b[1;32m     78\u001b[0m amax_with_dims \u001b[38;5;241m=\u001b[39m amax \u001b[38;5;28;01mif\u001b[39;00m keepdims \u001b[38;5;28;01melse\u001b[39;00m lax\u001b[38;5;241m.\u001b[39mexpand_dims(amax, pos_dims)\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/numpy/reductions.py:256\u001b[0m, in \u001b[0;36mmax\u001b[0;34m(a, axis, out, keepdims, initial, where)\u001b[0m\n\u001b[1;32m    252\u001b[0m \u001b[38;5;129m@_wraps\u001b[39m(np\u001b[38;5;241m.\u001b[39mmax, skip_params\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mout\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmax\u001b[39m(a: ArrayLike, axis: Axis \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, out: \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    254\u001b[0m         keepdims: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, initial: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    255\u001b[0m         where: Optional[ArrayLike] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[0;32m--> 256\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_reduce_max\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m_ensure_optional_axes\u001b[49m\u001b[43m(\u001b[49m\u001b[43maxis\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    257\u001b[0m \u001b[43m                     \u001b[49m\u001b[43mkeepdims\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkeepdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minitial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minitial\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwhere\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwhere\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 17 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/lib/python3.8/site-packages/jax/_src/interpreters/ad.py:314\u001b[0m, in \u001b[0;36mJVPTrace.process_primitive\u001b[0;34m(self, primitive, tracers, params)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m jvp:\n\u001b[1;32m    313\u001b[0m   msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDifferentiation rule for \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprimitive\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m not implemented\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m--> 314\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(msg)\n\u001b[1;32m    315\u001b[0m primal_out, tangent_out \u001b[38;5;241m=\u001b[39m jvp(primals_in, tangents_in, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams)\n\u001b[1;32m    316\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m primitive\u001b[38;5;241m.\u001b[39mmultiple_results:\n",
            "\u001b[0;31mNotImplementedError\u001b[0m: Differentiation rule for 'pmax' not implemented"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ZOcBSkNKEI6I"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}