Chapter-3/JAX_in_Action_Chapter_3_Array_with_TPU.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e6Pq2QZCNMYV",
        "outputId": "1735cf95-919c-45ac-eb56-e40f2f0f24ab"
      },
      "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]>=0.2.16\n",
            "  Downloading jax-0.4.13.tar.gz (1.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.3 MB 4.9 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 85.4 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 84.7 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.0 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 92.6 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]>=0.2.16) (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]>=0.2.16) (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 \bdone\n",
            "\u001b[?25h  Created wheel for jax: filename=jax-0.4.13-py3-none-any.whl size=1518704 sha256=7afb807cb668575c183cb10cdbd97c921d27b48a4c39c2a91517322970538231\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": [
        "import jax\n",
        "jax.__version__"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vcRv42-IQfZc",
        "outputId": "c2139eda-68c9-4f22-a2c4-18523f73c7cf"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.4.13'"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "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": "9JQtp7xPB0TS",
        "outputId": "fd451881-09db-4aba-b79f-283022bcf4fd"
      },
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZE6H-14Adw5d"
      },
      "source": [
        "## DeviceArray properties"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "Xw26fWWnf8eB"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tyk_72Oxq4WG",
        "outputId": "5f1306ea-e8ae-46ac-bca9-f19d699f964c"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([    1,    42, 31337])"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "np.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pZvJaXWTBm04",
        "outputId": "3f6248d2-c173-4033-c484-f908aa33deba"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JF5SSt3DBskh",
        "outputId": "ee839c15-8dbd-4227-f0e5-a29798c202f9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "31380"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "np.sum([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5O9F4CNYBymQ",
        "outputId": "5ebc56e9-f5da-4473-855b-1478c0633c53"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sum requires ndarray or scalar arguments, got <class 'list'> at position 0.\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  jnp.sum([1, 42, 31337])\n",
        "except TypeError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x76jmBRaHvQp",
        "outputId": "513b267b-748d-46b6-9e12-d6c80009d317"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(31380, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ],
      "source": [
        "jnp.sum(jnp.array([1, 42, 31337]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "sX1UevU0K5AW"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pDSN47PTRZHG",
        "outputId": "23d93fff-6682-4422-a854-e64f63ced51b"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ],
      "source": [
        "arr.ndim"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pghF6HXLx5O",
        "outputId": "0e5fee69-0f64-48e9-e8b3-fd99acb558b2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(3,)"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ],
      "source": [
        "arr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AXaNW2zaSzuR",
        "outputId": "17074f4e-ebdc-49bb-f6ea-618bfa28652f"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('int32')"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ],
      "source": [
        "arr.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G2GJkI3uRYYC",
        "outputId": "c26a12ee-3fc1-4456-ffe7-b71bfad36760"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "3"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ],
      "source": [
        "arr.size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rDB_80JCSv_b",
        "outputId": "a7d876c8-3696-4641-fd84-3751f28f9ef9"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "12"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "arr.nbytes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BXbIZE5Vd39Z"
      },
      "source": [
        "## Devices"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "1GAG9wYcaoSL"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gInETIcSakx6",
        "outputId": "bde64c52-53b2-4804-9a95-72e399d8c9e3"
      },
      "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": 16
        }
      ],
      "source": [
        "jax.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AUOges5na8X9",
        "outputId": "0dd1e52c-6d02-4627-cd32-14269aef7423"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[CpuDevice(id=0)]"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "jax.devices('cpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cXKsGEZObOkH",
        "outputId": "20b24d54-7374-46d3-b4cf-d3c5c93e61bd"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "8"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ],
      "source": [
        "jax.device_count('tpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LIFmPDWdamno",
        "outputId": "d1ab4061-48ff-4275-8435-25c552733ebb"
      },
      "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": 20
        }
      ],
      "source": [
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "sy08TlUVfgY_"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "beWtlcssfl1g",
        "outputId": "807c52d0-570b-4dab-d9e9-3c75f35bd817"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ],
      "source": [
        "arr.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "HQrwov7Zig0t"
      },
      "outputs": [],
      "source": [
        "arr_cpu = jax.device_put(arr, jax.devices('cpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fvhr-gSAivKy",
        "outputId": "14b28d89-7792-4289-c7d5-6864a7900646"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ],
      "source": [
        "arr_cpu.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "99hqLCwLin-w",
        "outputId": "179fa4e7-1c64-4530-d8df-cfd6a9888634"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ],
      "source": [
        "arr.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "9aQL3DFtjJvm"
      },
      "outputs": [],
      "source": [
        "arr_host = jax.device_get(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S6ep8EUPjarr",
        "outputId": "f2be1f0d-394f-447b-dbfb-6960547245d8"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "numpy.ndarray"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ],
      "source": [
        "type(arr_host)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7ela_gWXjRde",
        "outputId": "ec2d841e-c649-43d8-d2e2-1020488d9fe7"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "arr_host"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8H25pRzgkNnb",
        "outputId": "f05caa6d-d490-4d5b-96bb-9d6135572f9d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([    2,    84, 62674], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ],
      "source": [
        "arr + arr_cpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "zMkcgodHknpu"
      },
      "outputs": [],
      "source": [
        "arr_tpu = jax.device_put(arr, jax.devices('tpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbpuMMGHk1tK",
        "outputId": "9e822e59-c092-4ee2-94f7-be9a783578e2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)"
            ]
          },
          "metadata": {},
          "execution_count": 31
        }
      ],
      "source": [
        "arr_tpu.device()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xxTzC7Juk3vV",
        "outputId": "6ad142df-f968-4d07-b1fc-53549b5abe66"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Received incompatible devices for jitted computation. Got argument x1 of jax.numpy.add with shape int32[3] and device ids [0] on platform TPU and argument x2 of jax.numpy.add with shape int32[3] and device ids [0] on platform CPU\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  arr_tpu + arr_cpu\n",
        "except ValueError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MimQQxJVd8OT"
      },
      "source": [
        "## Asyncronous dispatch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "dkMYNMK8d-oZ"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "8TTzIPOud_c6"
      },
      "outputs": [],
      "source": [
        "a = jnp.array(range(1000000)).reshape((1000,1000))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WL4cK2EgeAOj",
        "outputId": "46dde9b9-1a3f-426f-926e-d5bbaa987e79"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1000, 1000)"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ],
      "source": [
        "a.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "81YXUMwfeUxL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "26ab8643-7136-45e8-a6ec-ffe928142150"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ],
      "source": [
        "a.device()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a,a)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xXmU97i2jB2u",
        "outputId": "6c43c4de-40dd-4b64-ff93-a42acbf2f355"
      },
      "execution_count": 55,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 1.07 ms, sys: 674 µs, total: 1.74 ms\n",
            "Wall time: 953 µs\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a,a).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q_LOBrs7jc1G",
        "outputId": "bdd6e049-5b02-4a75-d328-978cab29ead5"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 1.85 ms, sys: 1.17 ms, total: 3.02 ms\n",
            "Wall time: 2.07 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = np.asarray(jnp.dot(a,a))"
      ],
      "metadata": {
        "id": "9wtnSBNGADQC",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c666fb93-db67-4f56-f2e1-5339853b82c7"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 5.43 ms, sys: 638 µs, total: 6.07 ms\n",
            "Wall time: 3.77 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_cpu = jax.device_put(a, jax.devices('cpu')[0])"
      ],
      "metadata": {
        "id": "vHAlLz9Ck_1O"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_cpu.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "blmetLnalEOq",
        "outputId": "3940c143-dffe-4d4d-f5f0-073ff4072df4"
      },
      "execution_count": 59,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 59
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%time x = jnp.dot(a_cpu,a_cpu).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XSqQAxFRlQ4h",
        "outputId": "90b1007a-ba70-4e07-c352-e87bfd0aab83"
      },
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 341 ms, sys: 29.6 ms, total: 371 ms\n",
            "Wall time: 5.19 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Immutability"
      ],
      "metadata": {
        "id": "OzNEJzADTvlC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "oKO_DNwDUPSf"
      },
      "execution_count": 65,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = jnp.array(range(10))\n",
        "a_np  = np.array(range(10))"
      ],
      "metadata": {
        "id": "pHs4TJm1Ty7A"
      },
      "execution_count": 66,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Eplq9iRKHDM",
        "outputId": "e48f6ba6-6ed5-48b5-d63a-4b0021c3b4f4"
      },
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 67
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5], a_jnp[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GjtKMYxnOqKO",
        "outputId": "e343c63d-759e-4873-e9f8-4047543419a8"
      },
      "execution_count": 68,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, Array(5, dtype=int32))"
            ]
          },
          "metadata": {},
          "execution_count": 68
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5] = 100"
      ],
      "metadata": {
        "id": "9sZC4YCoUUnm"
      },
      "execution_count": 69,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_np[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HaRvmN0SUWnR",
        "outputId": "a43b7798-29b6-42d9-8807-bc040c1dff19"
      },
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "100"
            ]
          },
          "metadata": {},
          "execution_count": 70
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "  a_jnp[5] = 100\n",
        "except TypeError as e:\n",
        "  print(e)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zO8nRaogT3-N",
        "outputId": "f29820fa-1e28-4093-f97f-ec646806d341"
      },
      "execution_count": 71,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "'<class 'jaxlib.xla_extension.ArrayImpl'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[5].set(100)"
      ],
      "metadata": {
        "id": "rxAB1h8jT7tb"
      },
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp[5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BhSZ0WrfZhsR",
        "outputId": "da2cc802-ba1a-432a-c4ff-c8ba1f3ff743"
      },
      "execution_count": 73,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(100, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 73
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp[42]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4W91vzNOZmSd",
        "outputId": "eb73715e-cfb5-4275-a5f4-f5d713a7dc2a"
      },
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 74
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KmbSWZrNirZd",
        "outputId": "e3b8bc79-96ed-4a6a-b4cc-079343d81a8c"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 75
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='clip')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PFs6_qehi0Hk",
        "outputId": "6ecd1d9e-4942-41a4-f922-f464df1f69d2"
      },
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 76
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='drop')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1F1dy5q7jLW3",
        "outputId": "793368d7-ec02-4ef8-d04b-a8406fa78f67"
      },
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(-2147483648, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 77
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp.at[42].get(mode='fill', fill_value=-1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "waFQQytXjcMl",
        "outputId": "cb3556df-ac71-402d-be84-5d6e4a6a4283"
      },
      "execution_count": 78,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(-1, dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 78
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[42].set(100)\n",
        "a_jnp"
      ],
      "metadata": {
        "id": "6Q-IQhYAjdga",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dacd12e7-dbb1-45d9-b49b-0dda37190647"
      },
      "execution_count": 79,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8,   9], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 79
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "a_jnp = a_jnp.at[42].set(100, mode='clip')\n",
        "a_jnp"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pDDyLIxKkA2",
        "outputId": "20422d66-de29-4cac-b17b-c358e38a3e59"
      },
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8, 100], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 80
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Working with float64"
      ],
      "metadata": {
        "id": "ZQoM_CmGOvk5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# this only works on startup!\n",
        "from jax.config import config\n",
        "config.update(\"jax_enable_x64\", True)"
      ],
      "metadata": {
        "id": "Mvz9-cWsOxPC"
      },
      "execution_count": 81,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "metadata": {
        "id": "MtSHfhWyQYFZ"
      },
      "execution_count": 82,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# this will not work on TPU backend. Try using CPU or GPU.\n",
        "x = jnp.array(range(10), dtype=jnp.float64)\n",
        "x.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BsGsX08bKyFF",
        "outputId": "4ebaf38a-05d9-463a-892a-d47fb1ceaf73"
      },
      "execution_count": 83,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "metadata": {},
          "execution_count": 83
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Jpy24rpO8xb",
        "outputId": "d8db480d-cb98-4ed6-ae05-0f1091443b11"
      },
      "execution_count": 84,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)"
            ]
          },
          "metadata": {},
          "execution_count": 84
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xc = jax.device_put(x, jax.devices('cpu')[0])"
      ],
      "metadata": {
        "id": "5c5u5JTGQ7wP"
      },
      "execution_count": 85,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "xc.device()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3vK1jjnlRCHF",
        "outputId": "df69d2fb-eb83-4d06-89fc-d6081aec2cda"
      },
      "execution_count": 86,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "CpuDevice(id=0)"
            ]
          },
          "metadata": {},
          "execution_count": 86
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xc.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9bw68etaRHZD",
        "outputId": "c96ba462-c83c-48e5-be27-56284a1db0c9"
      },
      "execution_count": 87,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "metadata": {},
          "execution_count": 87
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16 = jnp.array(range(10), dtype=jnp.bfloat16)\n",
        "xb16.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DOFosWoRRJY0",
        "outputId": "22a63b83-f9d1-4c65-805a-27dc761c30ba"
      },
      "execution_count": 88,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype(bfloat16)"
            ]
          },
          "metadata": {},
          "execution_count": 88
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16.nbytes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jcY_9IgXSzty",
        "outputId": "2e3e00a1-d528-4c33-cee5-0d45a4408baa"
      },
      "execution_count": 89,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "metadata": {},
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x16 = jnp.array(range(10), dtype=jnp.float16)\n",
        "x16.dtype"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDxZWsldS4qU",
        "outputId": "06493891-012f-43c0-d566-7d843cfa7c22"
      },
      "execution_count": 90,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dtype('float16')"
            ]
          },
          "metadata": {},
          "execution_count": 90
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x16.nbytes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Z4NcluiUouv",
        "outputId": "a0b388c9-99f7-4b60-994b-52a39b6994b5"
      },
      "execution_count": 91,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "metadata": {},
          "execution_count": 91
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16+x16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jC6ynjHBUp4q",
        "outputId": "6ac49363-3074-4689-9eeb-e46181674edb"
      },
      "execution_count": 92,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 92
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "xb16+xb16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d6KaXrhYVget",
        "outputId": "078235df-c268-4220-b0fe-ee269e1bcfcf"
      },
      "execution_count": 93,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18], dtype=bfloat16)"
            ]
          },
          "metadata": {},
          "execution_count": 93
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## jax.numpy & jax.lax"
      ],
      "metadata": {
        "id": "XzqPR8GS-YzE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config.update(\"jax_enable_x64\", False)"
      ],
      "metadata": {
        "id": "D--HPBkMOsCB"
      },
      "execution_count": 94,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import lax\n",
        "from jax import random"
      ],
      "metadata": {
        "id": "kNemPL6xVjSK"
      },
      "execution_count": 95,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.add(42, 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RE4Smt9e-ixs",
        "outputId": "6435dc7b-91e8-4548-9620-ffe8cddcd725"
      },
      "execution_count": 96,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 96
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.add(42.0, 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z8GOhmXB-mTr",
        "outputId": "ea433f77-2efe-46b5-a3c2-42c8b0a520fd"
      },
      "execution_count": 97,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "metadata": {},
          "execution_count": 97
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "   lax.add(42, 42.0)\n",
        "except TypeError as e:\n",
        "  print(e)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G1JzxdU_-rDS",
        "outputId": "267370d8-9ab6-40b3-bcac-4de3b7d3f111"
      },
      "execution_count": 98,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "lax.add requires arguments to have the same dtypes, got int32, float32. (Tip: jnp.add is a similar function that does automatic type promotion on inputs).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "lax.add(jnp.float32(42), 42.0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MTBx3GE6-uNs",
        "outputId": "15615d7b-9fbf-4220-dbcb-ffcac0a8fe2c"
      },
      "execution_count": 99,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(84., dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def random_augmentation(image, augmentations, rng_key):\n",
        "   '''A function that applies a random transformation to an image'''\n",
        "   augmentation_index = random.randint(key=rng_key, minval=0, maxval=len(augmentations), shape=())\n",
        "   augmented_image = lax.switch(augmentation_index, augmentations, image)\n",
        "   return augmented_image"
      ],
      "metadata": {
        "id": "7fyJ9edV-8qO"
      },
      "execution_count": 100,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "add_noise_func = lambda x: x+10\n",
        "horizontal_flip_func = lambda x: x+1\n",
        "rotate_func = lambda x: x+2\n",
        "adjust_colors_func = lambda x: x+3\n",
        "\n",
        "augmentations = [\n",
        "   add_noise_func,\n",
        "   horizontal_flip_func,\n",
        "   rotate_func,\n",
        "   adjust_colors_func\n",
        "]\n"
      ],
      "metadata": {
        "id": "3WJ66AgB_D_o"
      },
      "execution_count": 101,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image = jnp.array(range(100))"
      ],
      "metadata": {
        "id": "aH-ZJ-DV_Vlu"
      },
      "execution_count": 102,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "random_augmentation(image, augmentations, random.PRNGKey(211))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oRDNt3q1_RJ4",
        "outputId": "95d0929c-ffae-46cc-aa97-b8eb7038e4ae"
      },
      "execution_count": 103,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([  1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,\n",
              "        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,\n",
              "        27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,\n",
              "        40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,\n",
              "        53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,\n",
              "        66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,\n",
              "        79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,\n",
              "        92,  93,  94,  95,  96,  97,  98,  99, 100], dtype=int32)"
            ]
          },
          "metadata": {},
          "execution_count": 103
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "mRahQ4xb_8ma"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}