Chapter-3/JAX_in_Action_Chapter_3_Array.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "vcRv42-IQfZc",
        "outputId": "c3acc10d-022e-49ed-b257-da39503bee03"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'0.4.30'"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import jax\n",
        "jax.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZE6H-14Adw5d"
      },
      "source": [
        "## DeviceArray properties"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Xw26fWWnf8eB"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tyk_72Oxq4WG",
        "outputId": "dbd7570a-f26a-4846-84da-546455e33792"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([    1,    42, 31337])"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pZvJaXWTBm04",
        "outputId": "9815ad1e-2231-4358-b9c6-cdc31f7101f2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JF5SSt3DBskh",
        "outputId": "5725bcf7-3b5f-4cad-aefe-ea5d1991011c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "31380"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.sum([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5O9F4CNYBymQ",
        "outputId": "ad00b7ee-b32e-41b7-eae6-b4bdc92b574a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "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": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x76jmBRaHvQp",
        "outputId": "1e2e8231-6e1c-4814-95a9-cd4e78343eae"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(31380, dtype=int32)"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.sum(jnp.array([1, 42, 31337]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "sX1UevU0K5AW"
      },
      "outputs": [],
      "source": [
        "arr = jnp.array([1, 42, 31337])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pDSN47PTRZHG",
        "outputId": "68594806-9c95-4870-ad93-817cf1302330"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.ndim"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pghF6HXLx5O",
        "outputId": "351b355a-4346-4649-fbee-a69d6fc3e5e1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(3,)"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AXaNW2zaSzuR",
        "outputId": "dda018c2-f981-482f-bb12-0c3119eb17f1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "dtype('int32')"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G2GJkI3uRYYC",
        "outputId": "55a064a5-9d87-4283-dc08-889404579800"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "3"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rDB_80JCSv_b",
        "outputId": "8055e6d3-e3f5-43f1-aa9f-eef3abe2b5e9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "12"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.nbytes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BXbIZE5Vd39Z"
      },
      "source": [
        "## Devices"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "1GAG9wYcaoSL"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gInETIcSakx6",
        "outputId": "253f0974-1d28-43c0-b7cb-0929d09f9089"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[cuda(id=0)]"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AUOges5na8X9",
        "outputId": "3a2208c7-8125-4940-a5a1-41e6a9a101c1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[CpuDevice(id=0)]"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.devices('cpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cXKsGEZObOkH",
        "outputId": "6eb84e74-300e-4425-d67f-669f75230974"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jax.device_count('gpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LIFmPDWdamno",
        "outputId": "8e21503c-4493-40a1-e68a-b710a4628b19"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[cuda(id=0)]"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "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": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "beWtlcssfl1g",
        "outputId": "c11ffe3e-7ee2-4bc0-f011-8c05ec81a2b7"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{cuda(id=0)}"
            ]
          },
          "execution_count": 23,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "HQrwov7Zig0t"
      },
      "outputs": [],
      "source": [
        "arr_cpu = jax.device_put(arr, jax.devices('cpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fvhr-gSAivKy",
        "outputId": "82622f22-7f7f-4007-b587-1c258e2964c0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{CpuDevice(id=0)}"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr_cpu.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "99hqLCwLin-w",
        "outputId": "d25ae8fe-eeae-4092-ffab-9e7b091b437a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{cuda(id=0)}"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "9aQL3DFtjJvm"
      },
      "outputs": [],
      "source": [
        "arr_host = jax.device_get(arr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "S6ep8EUPjarr",
        "outputId": "94d97bf9-181b-4355-b9a9-7669858a1a28"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "numpy.ndarray"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(arr_host)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7ela_gWXjRde",
        "outputId": "b6f35852-7fc4-4254-9bb1-9e789ead6f89"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([    1,    42, 31337], dtype=int32)"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr_host"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8H25pRzgkNnb",
        "outputId": "7e9f7939-2a9e-405f-d89b-16f809e90899"
      },
      "outputs": [],
      "source": [
        "arr2 = arr + arr_cpu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{CpuDevice(id=0)}"
            ]
          },
          "execution_count": 40,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr2.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "zMkcgodHknpu"
      },
      "outputs": [],
      "source": [
        "arr_gpu = jax.device_put(arr, jax.devices('gpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gbpuMMGHk1tK",
        "outputId": "afd80fce-0c22-40a8-d7e2-8ee2392bb3c2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{cuda(id=0)}"
            ]
          },
          "execution_count": 34,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "arr_gpu.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xxTzC7Juk3vV",
        "outputId": "4d0636a2-455b-41f3-900d-90749214f601"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Received incompatible devices for jitted computation. Got argument x1 of jax.numpy.add with shape int32[3] and device ids [0] on platform GPU and argument x2 of jax.numpy.add with shape int32[3] and device ids [0] on platform CPU\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  arr_gpu + arr_cpu\n",
        "except ValueError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MimQQxJVd8OT"
      },
      "source": [
        "## Asyncronous dispatch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "dkMYNMK8d-oZ"
      },
      "outputs": [],
      "source": [
        "import jax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "8TTzIPOud_c6"
      },
      "outputs": [],
      "source": [
        "a = jnp.array(range(1000000)).reshape((1000,1000))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WL4cK2EgeAOj",
        "outputId": "25e520dd-30ef-41de-9391-a38670ba538a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(1000, 1000)"
            ]
          },
          "execution_count": 52,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "81YXUMwfeUxL",
        "outputId": "6dcb7482-2220-4fde-dd61-03d2661f8898"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{cuda(id=0)}"
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xXmU97i2jB2u",
        "outputId": "93bb84fc-c608-4d64-f6ac-f9e1161ed041"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CPU times: user 656 μs, sys: 0 ns, total: 656 μs\n",
            "Wall time: 404 μs\n"
          ]
        }
      ],
      "source": [
        "%time x = jnp.dot(a,a)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q_LOBrs7jc1G",
        "outputId": "269a51f0-1476-4f26-e3a8-a221318c7d19"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CPU times: user 886 μs, sys: 719 μs, total: 1.61 ms\n",
            "Wall time: 1.7 ms\n"
          ]
        }
      ],
      "source": [
        "%time x = jnp.dot(a,a).block_until_ready()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9wtnSBNGADQC",
        "outputId": "f1aa607e-4a71-4570-d2e4-c330b7119e43"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CPU times: user 870 μs, sys: 706 μs, total: 1.58 ms\n",
            "Wall time: 2.14 ms\n"
          ]
        }
      ],
      "source": [
        "%time x = np.asarray(jnp.dot(a,a))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "id": "vHAlLz9Ck_1O"
      },
      "outputs": [],
      "source": [
        "a_cpu = jax.device_put(a, jax.devices('cpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "blmetLnalEOq",
        "outputId": "349c8817-96c7-4f09-e8cd-71d7b4d94495"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{CpuDevice(id=0)}"
            ]
          },
          "execution_count": 76,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_cpu.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XSqQAxFRlQ4h",
        "outputId": "9b4d98a6-71f1-4529-9b9f-ef58ecce5686"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CPU times: user 111 ms, sys: 0 ns, total: 111 ms\n",
            "Wall time: 9.27 ms\n"
          ]
        }
      ],
      "source": [
        "%time x = jnp.dot(a_cpu,a_cpu).block_until_ready()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mhY8ShWHeBIc"
      },
      "source": [
        "## GPU diagnostics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VnA4LuxbTyKb",
        "outputId": "8f2dd07e-ebc6-4dc8-8d18-c68a1d5d7736"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "nvcc: NVIDIA (R) Cuda compiler driver\n",
            "Copyright (c) 2005-2024 NVIDIA Corporation\n",
            "Built on Thu_Mar_28_02:18:24_PDT_2024\n",
            "Cuda compilation tools, release 12.4, V12.4.131\n",
            "Build cuda_12.4.r12.4/compiler.34097967_0\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/jet/.pyenv/versions/3.11.0/lib/python3.11/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
            "  pid, fd = os.forkpty()\n"
          ]
        }
      ],
      "source": [
        "!nvcc --version"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "id": "DtPcc6kJUE50"
      },
      "outputs": [],
      "source": [
        "#!pip install --upgrade \"jax[cuda]\" -f https://storage.googleapis.com/jax-releases/jax_releases.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sI5-pu0SUN_n",
        "outputId": "daea3000-4b8f-4c42-9674-e37f77623595"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sun Aug  4 20:59:49 2024       \n",
            "+-----------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 560.27                 Driver Version: 560.70         CUDA Version: 12.6     |\n",
            "|-----------------------------------------+------------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                        |               MIG M. |\n",
            "|=========================================+========================+======================|\n",
            "|   0  NVIDIA GeForce RTX 4070 ...    On  |   00000000:01:00.0  On |                  N/A |\n",
            "|  0%   40C    P8              8W /  285W |   15787MiB /  16376MiB |      5%      Default |\n",
            "|                                         |                        |                  N/A |\n",
            "+-----------------------------------------+------------------------+----------------------+\n",
            "                                                                                         \n",
            "+-----------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                              |\n",
            "|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |\n",
            "|        ID   ID                                                               Usage      |\n",
            "|=========================================================================================|\n",
            "|    0   N/A  N/A     25006      C   /python3.11                                 N/A      |\n",
            "+-----------------------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n",
        "  print('and then re-execute this cell.')\n",
        "else:\n",
        "  print(gpu_info)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OzNEJzADTvlC"
      },
      "source": [
        "## Immutability"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "id": "oKO_DNwDUPSf"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 82,
      "metadata": {
        "id": "pHs4TJm1Ty7A"
      },
      "outputs": [],
      "source": [
        "a_jnp = jnp.array(range(10))\n",
        "a_np  = np.array(range(10))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Eplq9iRKHDM",
        "outputId": "3ac48a01-918f-4467-ddea-bb720cdd7882"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
            ]
          },
          "execution_count": 83,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GjtKMYxnOqKO",
        "outputId": "4d089b95-00ba-4453-b2ac-e5aea0733c8d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(5, Array(5, dtype=int32))"
            ]
          },
          "execution_count": 84,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_np[5], a_jnp[5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 85,
      "metadata": {
        "id": "9sZC4YCoUUnm"
      },
      "outputs": [],
      "source": [
        "a_np[5] = 100"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HaRvmN0SUWnR",
        "outputId": "561ac73b-40b2-4e64-cb4c-e419d54b2a48"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "100"
            ]
          },
          "execution_count": 86,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_np[5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 87,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zO8nRaogT3-N",
        "outputId": "423b3fc2-400f-43f5-84d0-5243b60e056e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "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"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  a_jnp[5] = 100\n",
        "except TypeError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "id": "rxAB1h8jT7tb"
      },
      "outputs": [],
      "source": [
        "a_jnp = a_jnp.at[5].set(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BhSZ0WrfZhsR",
        "outputId": "6fac9383-8996-4098-85cc-8950f52312ab"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(100, dtype=int32)"
            ]
          },
          "execution_count": 89,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp[5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4W91vzNOZmSd",
        "outputId": "e334f000-ceca-4aa0-f2a9-424eb47c2f4d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "execution_count": 90,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp[42]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KmbSWZrNirZd",
        "outputId": "352768a7-77be-4b30-bbfb-7ae97b2ecefc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "execution_count": 91,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp.at[42].get()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PFs6_qehi0Hk",
        "outputId": "6d6926e9-aa6c-4854-fe1c-a164486e408f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(9, dtype=int32)"
            ]
          },
          "execution_count": 92,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp.at[42].get(mode='clip')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1F1dy5q7jLW3",
        "outputId": "d7f13888-caaf-4edd-f94d-603cdc181fab"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(-2147483648, dtype=int32)"
            ]
          },
          "execution_count": 93,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp.at[42].get(mode='drop')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "waFQQytXjcMl",
        "outputId": "064d7a9b-008f-4c92-dab6-a9705900b744"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(-1, dtype=int32)"
            ]
          },
          "execution_count": 94,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp.at[42].get(mode='fill', fill_value=-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Q-IQhYAjdga",
        "outputId": "3e2ab981-6cf2-4c19-d3ad-4c99d0c607d1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8,   9], dtype=int32)"
            ]
          },
          "execution_count": 95,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp = a_jnp.at[42].set(100)\n",
        "a_jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6pDDyLIxKkA2",
        "outputId": "0634a817-7c36-4bab-93d9-26475d95b369"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([  0,   1,   2,   3,   4, 100,   6,   7,   8, 100], dtype=int32)"
            ]
          },
          "execution_count": 96,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "a_jnp = a_jnp.at[42].set(100, mode='clip')\n",
        "a_jnp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZQoM_CmGOvk5"
      },
      "source": [
        "## Working with float64"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "id": "Mvz9-cWsOxPC"
      },
      "outputs": [],
      "source": [
        "# this only works on startup!\n",
        "import jax\n",
        "jax.config.update(\"jax_enable_x64\", True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "id": "MtSHfhWyQYFZ"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BsGsX08bKyFF",
        "outputId": "3660e173-3681-46b0-9a47-0b6bd0e1f7d6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "execution_count": 100,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 102,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Jpy24rpO8xb",
        "outputId": "3865f47b-a6a9-467b-f7e0-5fc433e30030"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{cuda(id=0)}"
            ]
          },
          "execution_count": 102,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 103,
      "metadata": {
        "id": "5c5u5JTGQ7wP"
      },
      "outputs": [],
      "source": [
        "xc = jax.device_put(x, jax.devices('cpu')[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3vK1jjnlRCHF",
        "outputId": "14e23c6d-8fea-473f-ada8-5a8040fc4136"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{CpuDevice(id=0)}"
            ]
          },
          "execution_count": 104,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xc.devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 105,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9bw68etaRHZD",
        "outputId": "f5b8dc1a-dbb8-402a-ae3e-a95a3d2cd038"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "dtype('float64')"
            ]
          },
          "execution_count": 105,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xc.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 106,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DOFosWoRRJY0",
        "outputId": "679ae9e6-b8b3-4977-8719-c778a6bd1ba2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "dtype(bfloat16)"
            ]
          },
          "execution_count": 106,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xb16 = jnp.array(range(10), dtype=jnp.bfloat16)\n",
        "xb16.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 107,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jcY_9IgXSzty",
        "outputId": "45f96941-20c0-4afe-a26b-7086e852f759"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "execution_count": 107,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xb16.nbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 108,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EDxZWsldS4qU",
        "outputId": "d881b0e3-9132-4226-f60a-772a1aec8d70"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "dtype('float16')"
            ]
          },
          "execution_count": 108,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x16 = jnp.array(range(10), dtype=jnp.float16)\n",
        "x16.dtype"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 109,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_Z4NcluiUouv",
        "outputId": "16348c97-2371-410d-f8d3-ce65fff32dc5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "execution_count": 109,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x16.nbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 110,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jC6ynjHBUp4q",
        "outputId": "083b4bd8-c5c8-4ae0-a349-457071922c48"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18.], dtype=float32)"
            ]
          },
          "execution_count": 110,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xb16+x16"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 111,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d6KaXrhYVget",
        "outputId": "e2317558-83b4-4ff2-f123-19cd67bf1740"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array([0, 2, 4, 6, 8, 10, 12, 14, 16, 18], dtype=bfloat16)"
            ]
          },
          "execution_count": 111,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "xb16+xb16"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XzqPR8GS-YzE"
      },
      "source": [
        "## jax.numpy & jax.lax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 112,
      "metadata": {
        "id": "KdN4dF-mf6bR"
      },
      "outputs": [],
      "source": [
        "jax.config.update(\"jax_enable_x64\", False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 113,
      "metadata": {
        "id": "kNemPL6xVjSK"
      },
      "outputs": [],
      "source": [
        "import jax.numpy as jnp\n",
        "from jax import lax\n",
        "from jax import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 114,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RE4Smt9e-ixs",
        "outputId": "78544f09-fd5c-463e-8263-812ec02ae0dc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "execution_count": 114,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.add(42, 42.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 115,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z8GOhmXB-mTr",
        "outputId": "7c4cf3f9-614e-41de-e6cf-da1ae9786575"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(84., dtype=float32, weak_type=True)"
            ]
          },
          "execution_count": 115,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "jnp.add(42.0, 42.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 116,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G1JzxdU_-rDS",
        "outputId": "7182bc37-7caa-4266-adbc-909ee3224ad2"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "Cannot lower jaxpr with verifier errors:\n\top requires the same element type for all operands and results\n\t\tat loc(\"jit(add)/jit(main)/add\"(callsite(\"<module>\"(\"/tmp/ipykernel_37180/1506868229.py\":2:3) at callsite(\"InteractiveShell.run_code\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3577:20) at callsite(\"InteractiveShell.run_ast_nodes\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3517:19) at callsite(\"InteractiveShell.run_cell_async\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3334:29) at callsite(\"_pseudo_sync_runner\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/async_helpers.py\":128:8) at callsite(\"InteractiveShell._run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3130:21) at callsite(\"InteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3075:21) at callsite(\"ZMQInteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/zmqshell.py\":549:15) at callsite(\"IPythonKernel.do_execute\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/ipkernel.py\":449:26) at \"Kernel.execute_request\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/kernelbase.py\":778:28))))))))))))\nDefine JAX_DUMP_IR_TO to dump the module.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mMLIRError\u001b[0m                                 Traceback (most recent call last)",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:970\u001b[0m, in \u001b[0;36mlower_jaxpr_to_module\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    969\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 970\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[43mctx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moperation\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mverify\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m    971\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    972\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot lower jaxpr with verifier errors. \u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m\n\u001b[1;32m    973\u001b[0m         dump_module_message(ctx\u001b[38;5;241m.\u001b[39mmodule, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mverification\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n",
            "\u001b[0;31mMLIRError\u001b[0m: Verification failed:\nerror: \"jit(add)/jit(main)/add\"(callsite(\"<module>\"(\"/tmp/ipykernel_37180/1506868229.py\":2:3) at callsite(\"InteractiveShell.run_code\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3577:20) at callsite(\"InteractiveShell.run_ast_nodes\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3517:19) at callsite(\"InteractiveShell.run_cell_async\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3334:29) at callsite(\"_pseudo_sync_runner\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/async_helpers.py\":128:8) at callsite(\"InteractiveShell._run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3130:21) at callsite(\"InteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3075:21) at callsite(\"ZMQInteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/zmqshell.py\":549:15) at callsite(\"IPythonKernel.do_execute\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/ipkernel.py\":449:26) at \"Kernel.execute_request\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/kernelbase.py\":778:28))))))))))): op requires the same element type for all operands and results",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[116], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m----> 2\u001b[0m    \u001b[43mlax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m42\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m42.0\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\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 e:\n\u001b[1;32m      4\u001b[0m   \u001b[38;5;28mprint\u001b[39m(e)\n",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/lax/lax.py:420\u001b[0m, in \u001b[0;36madd\u001b[0;34m(x, y)\u001b[0m\n\u001b[1;32m    418\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21madd\u001b[39m(x: ArrayLike, y: ArrayLike) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Array:\n\u001b[1;32m    419\u001b[0m \u001b[38;5;250m  \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Elementwise addition: :math:`x + y`.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 420\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43madd_p\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\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",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/core.py:416\u001b[0m, in \u001b[0;36mPrimitive.bind\u001b[0;34m(self, *args, **params)\u001b[0m\n\u001b[1;32m    413\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[1;32m    414\u001b[0m   \u001b[38;5;28;01massert\u001b[39;00m (\u001b[38;5;129;01mnot\u001b[39;00m config\u001b[38;5;241m.\u001b[39menable_checks\u001b[38;5;241m.\u001b[39mvalue \u001b[38;5;129;01mor\u001b[39;00m\n\u001b[1;32m    415\u001b[0m           \u001b[38;5;28mall\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(arg, Tracer) \u001b[38;5;129;01mor\u001b[39;00m valid_jaxtype(arg) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args)), args\n\u001b[0;32m--> 416\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind_with_trace\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfind_top_trace\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/core.py:420\u001b[0m, in \u001b[0;36mPrimitive.bind_with_trace\u001b[0;34m(self, trace, args, params)\u001b[0m\n\u001b[1;32m    418\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind_with_trace\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, args, params):\n\u001b[1;32m    419\u001b[0m   \u001b[38;5;28;01mwith\u001b[39;00m pop_level(trace\u001b[38;5;241m.\u001b[39mlevel):\n\u001b[0;32m--> 420\u001b[0m     out \u001b[38;5;241m=\u001b[39m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_primitive\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mmap\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfull_raise\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    421\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, out) \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmultiple_results \u001b[38;5;28;01melse\u001b[39;00m full_lower(out)\n",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/core.py:921\u001b[0m, in \u001b[0;36mEvalTrace.process_primitive\u001b[0;34m(self, primitive, tracers, params)\u001b[0m\n\u001b[1;32m    919\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m call_impl_with_key_reuse_checks(primitive, primitive\u001b[38;5;241m.\u001b[39mimpl, \u001b[38;5;241m*\u001b[39mtracers, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams)\n\u001b[1;32m    920\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 921\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[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/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/dispatch.py:87\u001b[0m, in \u001b[0;36mapply_primitive\u001b[0;34m(prim, *args, **params)\u001b[0m\n\u001b[1;32m     85\u001b[0m prev \u001b[38;5;241m=\u001b[39m lib\u001b[38;5;241m.\u001b[39mjax_jit\u001b[38;5;241m.\u001b[39mswap_thread_local_state_disable_jit(\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 87\u001b[0m   outs \u001b[38;5;241m=\u001b[39m \u001b[43mfun\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     88\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     89\u001b[0m   lib\u001b[38;5;241m.\u001b[39mjax_jit\u001b[38;5;241m.\u001b[39mswap_thread_local_state_disable_jit(prev)\n",
            "    \u001b[0;31m[... skipping hidden 15 frame]\u001b[0m\n",
            "File \u001b[0;32m~/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:983\u001b[0m, in \u001b[0;36mlower_jaxpr_to_module\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    981\u001b[0m   \u001b[38;5;28;01mfor\u001b[39;00m d \u001b[38;5;129;01min\u001b[39;00m e\u001b[38;5;241m.\u001b[39merror_diagnostics:\n\u001b[1;32m    982\u001b[0m     emit_diagnostic_info(d)\n\u001b[0;32m--> 983\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(msg_lines) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m\n\u001b[1;32m    984\u001b[0m                    dump_module_message(ctx\u001b[38;5;241m.\u001b[39mmodule, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mverification\u001b[39m\u001b[38;5;124m\"\u001b[39m)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m    986\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m LoweringResult(ctx\u001b[38;5;241m.\u001b[39mmodule, ctx\u001b[38;5;241m.\u001b[39mkeepalives, ctx\u001b[38;5;241m.\u001b[39mhost_callbacks,\n\u001b[1;32m    987\u001b[0m                       ctx\u001b[38;5;241m.\u001b[39mshape_poly_state)\n",
            "\u001b[0;31mValueError\u001b[0m: Cannot lower jaxpr with verifier errors:\n\top requires the same element type for all operands and results\n\t\tat loc(\"jit(add)/jit(main)/add\"(callsite(\"<module>\"(\"/tmp/ipykernel_37180/1506868229.py\":2:3) at callsite(\"InteractiveShell.run_code\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3577:20) at callsite(\"InteractiveShell.run_ast_nodes\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3517:19) at callsite(\"InteractiveShell.run_cell_async\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3334:29) at callsite(\"_pseudo_sync_runner\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/async_helpers.py\":128:8) at callsite(\"InteractiveShell._run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3130:21) at callsite(\"InteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/IPython/core/interactiveshell.py\":3075:21) at callsite(\"ZMQInteractiveShell.run_cell\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/zmqshell.py\":549:15) at callsite(\"IPythonKernel.do_execute\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/ipkernel.py\":449:26) at \"Kernel.execute_request\"(\"/home/jet/.local/share/virtualenvs/JAX-in-Action-2DrGcPDg/lib/python3.11/site-packages/ipykernel/kernelbase.py\":778:28))))))))))))\nDefine JAX_DUMP_IR_TO to dump the module."
          ]
        }
      ],
      "source": [
        "try:\n",
        "   lax.add(42, 42.0)\n",
        "except TypeError as e:\n",
        "  print(e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 117,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MTBx3GE6-uNs",
        "outputId": "bd28dd1e-650a-4173-eb06-ac223d9ede82"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Array(84., dtype=float32)"
            ]
          },
          "execution_count": 117,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "lax.add(jnp.float32(42), 42.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 118,
      "metadata": {
        "id": "7fyJ9edV-8qO"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 119,
      "metadata": {
        "id": "3WJ66AgB_D_o"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 120,
      "metadata": {
        "id": "aH-ZJ-DV_Vlu"
      },
      "outputs": [],
      "source": [
        "image = jnp.array(range(100))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 121,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oRDNt3q1_RJ4",
        "outputId": "0b594b1f-9f4d-4e11-aafb-7ed6a975e527"
      },
      "outputs": [
        {
          "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)"
            ]
          },
          "execution_count": 121,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "random_augmentation(image, augmentations, random.PRNGKey(211))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mRahQ4xb_8ma"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}