notebooks/all_jaxton_solutions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "826af5a8",
   "metadata": {
    "papermill": {
     "duration": 0.082175,
     "end_time": "2022-04-07T09:09:07.269203",
     "exception": false,
     "start_time": "2022-04-07T09:09:07.187028",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center>\n",
    "    <h1>JaxTon</h1>\n",
    "    <i>💯 JAX exercises</i>\n",
    "    <br>\n",
    "    <br>\n",
    "    <a href='https://github.com/vopani/jaxton/blob/master/LICENSE'>\n",
    "        <img src='https://img.shields.io/badge/license-Apache%202.0-blue.svg?logo=apache'>\n",
    "    </a>\n",
    "    <a href='https://github.com/vopani/jaxton'>\n",
    "        <img src='https://img.shields.io/github/stars/vopani/jaxton?color=yellowgreen&logo=github'>\n",
    "    </a>\n",
    "    <a href='https://twitter.com/vopani'>\n",
    "        <img src='https://img.shields.io/twitter/follow/vopani'>\n",
    "    </a>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7311d9",
   "metadata": {
    "papermill": {
     "duration": 0.079925,
     "end_time": "2022-04-07T09:09:07.431606",
     "exception": false,
     "start_time": "2022-04-07T09:09:07.351681",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center>\n",
    "    All solutions of <b>JaxTon</b>: <i>💯 JAX exercises</i>\n",
    "    <br>\n",
    "    You can find individual sets of the exercises and solutions on <a href=\"https://github.com/vopani/jaxton#exercises-\">GitHub</a>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20ebce84",
   "metadata": {
    "papermill": {
     "duration": 0.095786,
     "end_time": "2022-04-07T09:09:07.607648",
     "exception": false,
     "start_time": "2022-04-07T09:09:07.511862",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Prerequisites**\n",
    "\n",
    "* The configuration of jax should be updated as shown in the code snippet below in order to use TPUs.\n",
    "* A sample array `sample_data` will be used for some of the exercises.\n",
    "* Sample functions `square`, `cube` and `areas` will be used for some of the exercises."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c3de5aa2",
   "metadata": {
    "_kg_hide-input": true,
    "_kg_hide-output": true,
    "execution": {
     "iopub.execute_input": "2022-04-07T09:09:07.782527Z",
     "iopub.status.busy": "2022-04-07T09:09:07.781332Z",
     "iopub.status.idle": "2022-04-07T09:09:17.077386Z",
     "shell.execute_reply": "2022-04-07T09:09:17.076362Z",
     "shell.execute_reply.started": "2022-04-07T06:11:29.612884Z"
    },
    "papermill": {
     "duration": 9.379282,
     "end_time": "2022-04-07T09:09:17.077587",
     "exception": false,
     "start_time": "2022-04-07T09:09:07.698305",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (0.2.19)\r\n",
      "Requirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax) (3.3.0)\r\n",
      "Requirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\r\n",
      "Requirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\r\n",
      "Requirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from absl-py->jax) (1.15.0)\r\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\r\n"
     ]
    }
   ],
   "source": [
    "## install jax\n",
    "!python3 -m pip install jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e5febc48",
   "metadata": {
    "_kg_hide-input": false,
    "execution": {
     "iopub.execute_input": "2022-04-07T09:09:17.249204Z",
     "iopub.status.busy": "2022-04-07T09:09:17.244872Z",
     "iopub.status.idle": "2022-04-07T09:10:06.477438Z",
     "shell.execute_reply": "2022-04-07T09:10:06.478296Z",
     "shell.execute_reply.started": "2022-04-07T06:11:40.551295Z"
    },
    "papermill": {
     "duration": 49.317754,
     "end_time": "2022-04-07T09:10:06.478558",
     "exception": false,
     "start_time": "2022-04-07T09:09:17.160804",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
       " TpuDevice(id=1, process_index=0, coords=(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)]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## import packages\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import math\n",
    "import os\n",
    "import requests\n",
    "\n",
    "## setup JAX to use TPUs if available\n",
    "try:\n",
    "    url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'\n",
    "    resp = requests.post(url)\n",
    "    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'\n",
    "    jax.config.FLAGS.jax_backend_target = os.environ['TPU_NAME']\n",
    "except:\n",
    "    pass\n",
    "\n",
    "jax.devices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "869e37b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:06.759062Z",
     "iopub.status.busy": "2022-04-07T09:10:06.758007Z",
     "iopub.status.idle": "2022-04-07T09:10:06.797502Z",
     "shell.execute_reply": "2022-04-07T09:10:06.798281Z",
     "shell.execute_reply.started": "2022-04-07T06:11:42.653081Z"
    },
    "papermill": {
     "duration": 0.185348,
     "end_time": "2022-04-07T09:10:06.798554",
     "exception": false,
     "start_time": "2022-04-07T09:10:06.613206",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([10,  1, 24, 20, 15, 14], dtype=int32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## sample data\n",
    "sample_data = jnp.array([10, 1, 24, 20, 15, 14])\n",
    "sample_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0228f3a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:07.078721Z",
     "iopub.status.busy": "2022-04-07T09:10:07.077552Z",
     "iopub.status.idle": "2022-04-07T09:10:07.081401Z",
     "shell.execute_reply": "2022-04-07T09:10:07.082295Z",
     "shell.execute_reply.started": "2022-04-07T06:11:42.669985Z"
    },
    "papermill": {
     "duration": 0.144474,
     "end_time": "2022-04-07T09:10:07.082692",
     "exception": false,
     "start_time": "2022-04-07T09:10:06.938218",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## sample square function\n",
    "def square(x):\n",
    "    return x**2\n",
    "\n",
    "square(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b2373a3b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:07.364205Z",
     "iopub.status.busy": "2022-04-07T09:10:07.363155Z",
     "iopub.status.idle": "2022-04-07T09:10:07.367603Z",
     "shell.execute_reply": "2022-04-07T09:10:07.368302Z",
     "shell.execute_reply.started": "2022-04-07T06:11:42.678633Z"
    },
    "papermill": {
     "duration": 0.150969,
     "end_time": "2022-04-07T09:10:07.368556",
     "exception": false,
     "start_time": "2022-04-07T09:10:07.217587",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## sample cube function\n",
    "def cube(x):\n",
    "    return x**3\n",
    "\n",
    "cube(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c6a76d2b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:07.640911Z",
     "iopub.status.busy": "2022-04-07T09:10:07.639871Z",
     "iopub.status.idle": "2022-04-07T09:10:07.643803Z",
     "shell.execute_reply": "2022-04-07T09:10:07.644539Z",
     "shell.execute_reply.started": "2022-04-07T06:11:42.692708Z"
    },
    "papermill": {
     "duration": 0.140812,
     "end_time": "2022-04-07T09:10:07.644795",
     "exception": false,
     "start_time": "2022-04-07T09:10:07.503983",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.7320508075688772, 4, 12.566370614359172]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## sample areas function\n",
    "def areas(x):\n",
    "    return [math.sqrt(3)*x**2/4, x**2, math.pi*x**2]\n",
    "\n",
    "areas(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ad3f5d",
   "metadata": {
    "papermill": {
     "duration": 0.132053,
     "end_time": "2022-04-07T09:10:07.915977",
     "exception": false,
     "start_time": "2022-04-07T09:10:07.783924",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 1: Install the `jax` package**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b4f079c0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:08.187877Z",
     "iopub.status.busy": "2022-04-07T09:10:08.187091Z",
     "iopub.status.idle": "2022-04-07T09:10:20.288760Z",
     "shell.execute_reply": "2022-04-07T09:10:20.289554Z",
     "shell.execute_reply.started": "2022-04-07T06:11:42.709432Z"
    },
    "papermill": {
     "duration": 12.244304,
     "end_time": "2022-04-07T09:10:20.289972",
     "exception": false,
     "start_time": "2022-04-07T09:10:08.045668",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (0.2.19)\r\n",
      "Requirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax) (3.3.0)\r\n",
      "Requirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\r\n",
      "Requirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\r\n",
      "Requirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from absl-py->jax) (1.15.0)\r\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\r\n"
     ]
    }
   ],
   "source": [
    "!python3 -m pip install jax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c4b51d",
   "metadata": {
    "papermill": {
     "duration": 0.132926,
     "end_time": "2022-04-07T09:10:20.559107",
     "exception": false,
     "start_time": "2022-04-07T09:10:20.426181",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 2: Import the `jax` package**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8de1dbf7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:20.837411Z",
     "iopub.status.busy": "2022-04-07T09:10:20.836492Z",
     "iopub.status.idle": "2022-04-07T09:10:20.838875Z",
     "shell.execute_reply": "2022-04-07T09:10:20.839606Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.805800Z"
    },
    "papermill": {
     "duration": 0.144939,
     "end_time": "2022-04-07T09:10:20.839826",
     "exception": false,
     "start_time": "2022-04-07T09:10:20.694887",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import jax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18af1225",
   "metadata": {
    "papermill": {
     "duration": 0.135049,
     "end_time": "2022-04-07T09:10:21.109445",
     "exception": false,
     "start_time": "2022-04-07T09:10:20.974396",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 3: Display the version of `jax`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "28b9430f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:21.376969Z",
     "iopub.status.busy": "2022-04-07T09:10:21.375985Z",
     "iopub.status.idle": "2022-04-07T09:10:21.379957Z",
     "shell.execute_reply": "2022-04-07T09:10:21.380703Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.814899Z"
    },
    "papermill": {
     "duration": 0.140957,
     "end_time": "2022-04-07T09:10:21.380934",
     "exception": false,
     "start_time": "2022-04-07T09:10:21.239977",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.2.19'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ec5d3d4",
   "metadata": {
    "papermill": {
     "duration": 0.136692,
     "end_time": "2022-04-07T09:10:21.651691",
     "exception": false,
     "start_time": "2022-04-07T09:10:21.514999",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 4: Display the default backend of `jax`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d82bce1a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:21.926275Z",
     "iopub.status.busy": "2022-04-07T09:10:21.924949Z",
     "iopub.status.idle": "2022-04-07T09:10:21.929096Z",
     "shell.execute_reply": "2022-04-07T09:10:21.929842Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.840449Z"
    },
    "papermill": {
     "duration": 0.144977,
     "end_time": "2022-04-07T09:10:21.930194",
     "exception": false,
     "start_time": "2022-04-07T09:10:21.785217",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'tpu'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.default_backend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09c8d54f",
   "metadata": {
    "papermill": {
     "duration": 0.138163,
     "end_time": "2022-04-07T09:10:22.198411",
     "exception": false,
     "start_time": "2022-04-07T09:10:22.060248",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 5: Display the devices of the backend**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cd3153c7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:22.482759Z",
     "iopub.status.busy": "2022-04-07T09:10:22.481737Z",
     "iopub.status.idle": "2022-04-07T09:10:22.486140Z",
     "shell.execute_reply": "2022-04-07T09:10:22.486877Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.851974Z"
    },
    "papermill": {
     "duration": 0.153248,
     "end_time": "2022-04-07T09:10:22.487112",
     "exception": false,
     "start_time": "2022-04-07T09:10:22.333864",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
       " TpuDevice(id=1, process_index=0, coords=(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)]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f9b2870",
   "metadata": {
    "papermill": {
     "duration": 0.13379,
     "end_time": "2022-04-07T09:10:22.763219",
     "exception": false,
     "start_time": "2022-04-07T09:10:22.629429",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 6: Create a JAX DeviceArray with values [10, 1, 24] and assign it to `data`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0d3a747f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:23.106920Z",
     "iopub.status.busy": "2022-04-07T09:10:23.105977Z",
     "iopub.status.idle": "2022-04-07T09:10:23.143763Z",
     "shell.execute_reply": "2022-04-07T09:10:23.143016Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.864738Z"
    },
    "papermill": {
     "duration": 0.179895,
     "end_time": "2022-04-07T09:10:23.143961",
     "exception": false,
     "start_time": "2022-04-07T09:10:22.964066",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([10,  1, 24], dtype=int32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = jax.numpy.array([10, 1, 24])\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7ae14dc",
   "metadata": {
    "papermill": {
     "duration": 0.14012,
     "end_time": "2022-04-07T09:10:23.421822",
     "exception": false,
     "start_time": "2022-04-07T09:10:23.281702",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 7: Display the type of `data`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2d940622",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:23.711803Z",
     "iopub.status.busy": "2022-04-07T09:10:23.710827Z",
     "iopub.status.idle": "2022-04-07T09:10:23.714726Z",
     "shell.execute_reply": "2022-04-07T09:10:23.715708Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.878615Z"
    },
    "papermill": {
     "duration": 0.153869,
     "end_time": "2022-04-07T09:10:23.716032",
     "exception": false,
     "start_time": "2022-04-07T09:10:23.562163",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "jax.interpreters.xla._DeviceArray"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57af9177",
   "metadata": {
    "papermill": {
     "duration": 0.138319,
     "end_time": "2022-04-07T09:10:23.989840",
     "exception": false,
     "start_time": "2022-04-07T09:10:23.851521",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 8: Display the shape of `data`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f9cc29f6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:24.277329Z",
     "iopub.status.busy": "2022-04-07T09:10:24.276365Z",
     "iopub.status.idle": "2022-04-07T09:10:24.281324Z",
     "shell.execute_reply": "2022-04-07T09:10:24.280580Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.889662Z"
    },
    "papermill": {
     "duration": 0.152703,
     "end_time": "2022-04-07T09:10:24.281527",
     "exception": false,
     "start_time": "2022-04-07T09:10:24.128824",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3,)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6efc0622",
   "metadata": {
    "papermill": {
     "duration": 0.138247,
     "end_time": "2022-04-07T09:10:24.559437",
     "exception": false,
     "start_time": "2022-04-07T09:10:24.421190",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 9: Transfer `data` to host and assign it to `data_host`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "80c11e0a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:24.860514Z",
     "iopub.status.busy": "2022-04-07T09:10:24.859488Z",
     "iopub.status.idle": "2022-04-07T09:10:24.864800Z",
     "shell.execute_reply": "2022-04-07T09:10:24.864012Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.902623Z"
    },
    "papermill": {
     "duration": 0.150861,
     "end_time": "2022-04-07T09:10:24.865007",
     "exception": false,
     "start_time": "2022-04-07T09:10:24.714146",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([10,  1, 24], dtype=int32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_host = jax.device_get(data)\n",
    "data_host"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c519a13",
   "metadata": {
    "papermill": {
     "duration": 0.140021,
     "end_time": "2022-04-07T09:10:25.137987",
     "exception": false,
     "start_time": "2022-04-07T09:10:24.997966",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 10: Transfer `data_host` to device and assign it to `data_device`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e3ce129b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:25.418994Z",
     "iopub.status.busy": "2022-04-07T09:10:25.418085Z",
     "iopub.status.idle": "2022-04-07T09:10:25.448453Z",
     "shell.execute_reply": "2022-04-07T09:10:25.449188Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.916553Z"
    },
    "papermill": {
     "duration": 0.172093,
     "end_time": "2022-04-07T09:10:25.449413",
     "exception": false,
     "start_time": "2022-04-07T09:10:25.277320",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([10,  1, 24], dtype=int32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_device = jax.device_put(data_host)\n",
    "data_device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5ad9140",
   "metadata": {
    "papermill": {
     "duration": 0.137924,
     "end_time": "2022-04-07T09:10:25.727448",
     "exception": false,
     "start_time": "2022-04-07T09:10:25.589524",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 11: Create a matrix with values [[10, 1, 24], [20, 15, 14]] and assign it to `data`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c19c9f96",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:26.014103Z",
     "iopub.status.busy": "2022-04-07T09:10:26.013117Z",
     "iopub.status.idle": "2022-04-07T09:10:26.051509Z",
     "shell.execute_reply": "2022-04-07T09:10:26.052278Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.930505Z"
    },
    "papermill": {
     "duration": 0.184884,
     "end_time": "2022-04-07T09:10:26.052541",
     "exception": false,
     "start_time": "2022-04-07T09:10:25.867657",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[10,  1, 24],\n",
       "             [20, 15, 14]], dtype=int32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = jnp.array([[10, 1, 24], [20, 15, 14]])\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50334883",
   "metadata": {
    "papermill": {
     "duration": 0.141699,
     "end_time": "2022-04-07T09:10:26.340517",
     "exception": false,
     "start_time": "2022-04-07T09:10:26.198818",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 12: Assign the transpose of `data` to `dataT`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "311a3aad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:10:26.653387Z",
     "iopub.status.busy": "2022-04-07T09:10:26.645671Z",
     "iopub.status.idle": "2022-04-07T09:11:28.532467Z",
     "shell.execute_reply": "2022-04-07T09:11:28.533297Z",
     "shell.execute_reply.started": "2022-04-07T06:11:51.943659Z"
    },
    "papermill": {
     "duration": 62.039199,
     "end_time": "2022-04-07T09:11:28.533560",
     "exception": false,
     "start_time": "2022-04-07T09:10:26.494361",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[10, 20],\n",
       "             [ 1, 15],\n",
       "             [24, 14]], dtype=int32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataT = data.T\n",
    "dataT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b9de8d3",
   "metadata": {
    "papermill": {
     "duration": 0.139674,
     "end_time": "2022-04-07T09:11:28.815380",
     "exception": false,
     "start_time": "2022-04-07T09:11:28.675706",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 13: Assign the element of `data` at index [0, 2] to `value`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cede4476",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:29.101214Z",
     "iopub.status.busy": "2022-04-07T09:11:29.100041Z",
     "iopub.status.idle": "2022-04-07T09:11:29.306130Z",
     "shell.execute_reply": "2022-04-07T09:11:29.305027Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.186169Z"
    },
    "papermill": {
     "duration": 0.353341,
     "end_time": "2022-04-07T09:11:29.306349",
     "exception": false,
     "start_time": "2022-04-07T09:11:28.953008",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(24, dtype=int32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "value = data[0, 2]\n",
    "value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5bc1a2",
   "metadata": {
    "papermill": {
     "duration": 0.144341,
     "end_time": "2022-04-07T09:11:29.592547",
     "exception": false,
     "start_time": "2022-04-07T09:11:29.448206",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 14: Update the value of `data` at index [1, 1] to `100`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "edd35ac7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:29.885787Z",
     "iopub.status.busy": "2022-04-07T09:11:29.884831Z",
     "iopub.status.idle": "2022-04-07T09:11:29.962676Z",
     "shell.execute_reply": "2022-04-07T09:11:29.961888Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.280195Z"
    },
    "papermill": {
     "duration": 0.226209,
     "end_time": "2022-04-07T09:11:29.962881",
     "exception": false,
     "start_time": "2022-04-07T09:11:29.736672",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 10,   1,  24],\n",
       "             [ 20, 100,  14]], dtype=int32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = data.at[1, 1].set(100)\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccc3c019",
   "metadata": {
    "papermill": {
     "duration": 0.147213,
     "end_time": "2022-04-07T09:11:30.253294",
     "exception": false,
     "start_time": "2022-04-07T09:11:30.106081",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 15: Add `41` to the value of `data` at index [0, 0]**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a1b62af2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:30.547980Z",
     "iopub.status.busy": "2022-04-07T09:11:30.546636Z",
     "iopub.status.idle": "2022-04-07T09:11:30.611854Z",
     "shell.execute_reply": "2022-04-07T09:11:30.613000Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.337873Z"
    },
    "papermill": {
     "duration": 0.215548,
     "end_time": "2022-04-07T09:11:30.613266",
     "exception": false,
     "start_time": "2022-04-07T09:11:30.397718",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 51,   1,  24],\n",
       "             [ 20, 100,  14]], dtype=int32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = data.at[0, 0].add(41)\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91b47a22",
   "metadata": {
    "papermill": {
     "duration": 0.144214,
     "end_time": "2022-04-07T09:11:30.908050",
     "exception": false,
     "start_time": "2022-04-07T09:11:30.763836",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 16: Calculate the minimum values over axis=1 and assign it to `mins`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d917fb2b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:31.202396Z",
     "iopub.status.busy": "2022-04-07T09:11:31.200647Z",
     "iopub.status.idle": "2022-04-07T09:11:31.230622Z",
     "shell.execute_reply": "2022-04-07T09:11:31.231573Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.373401Z"
    },
    "papermill": {
     "duration": 0.182741,
     "end_time": "2022-04-07T09:11:31.231832",
     "exception": false,
     "start_time": "2022-04-07T09:11:31.049091",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 1, 14], dtype=int32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mins = data.min(axis=1)\n",
    "mins"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9fb35fb",
   "metadata": {
    "papermill": {
     "duration": 0.143653,
     "end_time": "2022-04-07T09:11:31.518055",
     "exception": false,
     "start_time": "2022-04-07T09:11:31.374402",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 17: Select the first row of values of `data` and assign it to `data_select`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fa393bca",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:31.815796Z",
     "iopub.status.busy": "2022-04-07T09:11:31.814391Z",
     "iopub.status.idle": "2022-04-07T09:11:31.868257Z",
     "shell.execute_reply": "2022-04-07T09:11:31.869032Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.401583Z"
    },
    "papermill": {
     "duration": 0.202965,
     "end_time": "2022-04-07T09:11:31.869275",
     "exception": false,
     "start_time": "2022-04-07T09:11:31.666310",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([51,  1, 24], dtype=int32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_select = data[0]\n",
    "data_select"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559247ee",
   "metadata": {
    "papermill": {
     "duration": 0.150149,
     "end_time": "2022-04-07T09:11:32.162498",
     "exception": false,
     "start_time": "2022-04-07T09:11:32.012349",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 18: Append the row `data_select` to `data`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "be5f46dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:32.462014Z",
     "iopub.status.busy": "2022-04-07T09:11:32.461051Z",
     "iopub.status.idle": "2022-04-07T09:11:32.509562Z",
     "shell.execute_reply": "2022-04-07T09:11:32.510318Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.428651Z"
    },
    "papermill": {
     "duration": 0.197428,
     "end_time": "2022-04-07T09:11:32.510574",
     "exception": false,
     "start_time": "2022-04-07T09:11:32.313146",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 51,   1,  24],\n",
       "             [ 20, 100,  14],\n",
       "             [ 51,   1,  24]], dtype=int32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = jnp.vstack([data, data_select])\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "defe3f65",
   "metadata": {
    "papermill": {
     "duration": 0.147704,
     "end_time": "2022-04-07T09:11:32.815303",
     "exception": false,
     "start_time": "2022-04-07T09:11:32.667599",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 19: Multiply the matrices `data` and `dataT` and assign it to `data_prod`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "54dbfa7d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:33.123335Z",
     "iopub.status.busy": "2022-04-07T09:11:33.122358Z",
     "iopub.status.idle": "2022-04-07T09:11:33.151624Z",
     "shell.execute_reply": "2022-04-07T09:11:33.152354Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.462042Z"
    },
    "papermill": {
     "duration": 0.187404,
     "end_time": "2022-04-07T09:11:33.152607",
     "exception": false,
     "start_time": "2022-04-07T09:11:32.965203",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[1087, 1371],\n",
       "             [ 636, 2096],\n",
       "             [1087, 1371]], dtype=int32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_prod = jnp.dot(data, dataT)\n",
    "data_prod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8435de0f",
   "metadata": {
    "papermill": {
     "duration": 0.147763,
     "end_time": "2022-04-07T09:11:33.448132",
     "exception": false,
     "start_time": "2022-04-07T09:11:33.300369",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 20: Convert the dtype of `data_prod` to `float32`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2968cdb2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:33.749249Z",
     "iopub.status.busy": "2022-04-07T09:11:33.748197Z",
     "iopub.status.idle": "2022-04-07T09:11:33.777540Z",
     "shell.execute_reply": "2022-04-07T09:11:33.776373Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.491705Z"
    },
    "papermill": {
     "duration": 0.182352,
     "end_time": "2022-04-07T09:11:33.777763",
     "exception": false,
     "start_time": "2022-04-07T09:11:33.595411",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[1087., 1371.],\n",
       "             [ 636., 2096.],\n",
       "             [1087., 1371.]], dtype=float32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_prod = jnp.array(data_prod, dtype=jnp.float32)\n",
    "data_prod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ed1bd4",
   "metadata": {
    "papermill": {
     "duration": 0.15089,
     "end_time": "2022-04-07T09:11:34.072749",
     "exception": false,
     "start_time": "2022-04-07T09:11:33.921859",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 21: Create a pseudorandom number generator key with seed=100 and assign it to `key`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "6148cb9f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:34.378944Z",
     "iopub.status.busy": "2022-04-07T09:11:34.377754Z",
     "iopub.status.idle": "2022-04-07T09:11:34.503058Z",
     "shell.execute_reply": "2022-04-07T09:11:34.503849Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.528235Z"
    },
    "papermill": {
     "duration": 0.284145,
     "end_time": "2022-04-07T09:11:34.504104",
     "exception": false,
     "start_time": "2022-04-07T09:11:34.219959",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([  0, 100], dtype=uint32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key = jax.random.PRNGKey(100)\n",
    "key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bee44165",
   "metadata": {
    "papermill": {
     "duration": 0.149087,
     "end_time": "2022-04-07T09:11:34.797414",
     "exception": false,
     "start_time": "2022-04-07T09:11:34.648327",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 22: Create a subkey from `key` and assign it to `subkey`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "90d89ac8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:35.121088Z",
     "iopub.status.busy": "2022-04-07T09:11:35.099919Z",
     "iopub.status.idle": "2022-04-07T09:11:35.341950Z",
     "shell.execute_reply": "2022-04-07T09:11:35.342639Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.598730Z"
    },
    "papermill": {
     "duration": 0.398873,
     "end_time": "2022-04-07T09:11:35.343018",
     "exception": false,
     "start_time": "2022-04-07T09:11:34.944145",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3011861781, 1867493174], dtype=uint32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key, subkey = jax.random.split(key)\n",
    "subkey"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f13926f",
   "metadata": {
    "papermill": {
     "duration": 0.150238,
     "end_time": "2022-04-07T09:11:35.642180",
     "exception": false,
     "start_time": "2022-04-07T09:11:35.491942",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 23: Split `key` into seven subkeys `key_1`, `key_2`, `key_3`, `key_4`, `key_5`, `key_6` and `key_7`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "36505e77",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:35.988336Z",
     "iopub.status.busy": "2022-04-07T09:11:35.946941Z",
     "iopub.status.idle": "2022-04-07T09:11:36.190868Z",
     "shell.execute_reply": "2022-04-07T09:11:36.191629Z",
     "shell.execute_reply.started": "2022-04-07T06:11:52.810859Z"
    },
    "papermill": {
     "duration": 0.395405,
     "end_time": "2022-04-07T09:11:36.191856",
     "exception": false,
     "start_time": "2022-04-07T09:11:35.796451",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 402730500, 1595431526], dtype=uint32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_1, key_2, key_3, key_4, key_5, key_6, key_7 = jax.random.split(key, num=7)\n",
    "key_1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64b808e",
   "metadata": {
    "papermill": {
     "duration": 0.145864,
     "end_time": "2022-04-07T09:11:36.483197",
     "exception": false,
     "start_time": "2022-04-07T09:11:36.337333",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 24: Create a random permutation of `sample_data` using `key_1` and assign it to `data_permutation`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f6330a61",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:36.804825Z",
     "iopub.status.busy": "2022-04-07T09:11:36.803791Z",
     "iopub.status.idle": "2022-04-07T09:11:37.262623Z",
     "shell.execute_reply": "2022-04-07T09:11:37.261773Z",
     "shell.execute_reply.started": "2022-04-07T06:11:53.054481Z"
    },
    "papermill": {
     "duration": 0.627623,
     "end_time": "2022-04-07T09:11:37.262807",
     "exception": false,
     "start_time": "2022-04-07T09:11:36.635184",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([20, 14,  1, 10, 24, 15], dtype=int32)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_permutation = jax.random.permutation(key_1, sample_data)\n",
    "data_permutation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51c70957",
   "metadata": {
    "papermill": {
     "duration": 0.143434,
     "end_time": "2022-04-07T09:11:37.547414",
     "exception": false,
     "start_time": "2022-04-07T09:11:37.403980",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 25: Choose a random element from `sample_data` using `key_2` and assign it to `random_selection`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "e37c5809",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:37.873734Z",
     "iopub.status.busy": "2022-04-07T09:11:37.872635Z",
     "iopub.status.idle": "2022-04-07T09:11:38.736356Z",
     "shell.execute_reply": "2022-04-07T09:11:38.737162Z",
     "shell.execute_reply.started": "2022-04-07T06:11:53.439418Z"
    },
    "papermill": {
     "duration": 1.044776,
     "end_time": "2022-04-07T09:11:38.737404",
     "exception": false,
     "start_time": "2022-04-07T09:11:37.692628",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(1, dtype=int32)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_selection = jax.random.choice(key_2, sample_data)\n",
    "random_selection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4aca1c1",
   "metadata": {
    "papermill": {
     "duration": 0.145518,
     "end_time": "2022-04-07T09:11:39.032325",
     "exception": false,
     "start_time": "2022-04-07T09:11:38.886807",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 26: Sample an integer between 10 and 24 using `key_3` and assign it to `sample_int`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "50333b29",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:39.366992Z",
     "iopub.status.busy": "2022-04-07T09:11:39.326653Z",
     "iopub.status.idle": "2022-04-07T09:11:40.306787Z",
     "shell.execute_reply": "2022-04-07T09:11:40.307556Z",
     "shell.execute_reply.started": "2022-04-07T06:11:53.919894Z"
    },
    "papermill": {
     "duration": 1.132729,
     "end_time": "2022-04-07T09:11:40.307792",
     "exception": false,
     "start_time": "2022-04-07T09:11:39.175063",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([14], dtype=int32)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_int = jax.random.randint(key_3, shape=(1,), minval=10, maxval=24)\n",
    "sample_int"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f7c6ffb",
   "metadata": {
    "papermill": {
     "duration": 0.147749,
     "end_time": "2022-04-07T09:11:40.607075",
     "exception": false,
     "start_time": "2022-04-07T09:11:40.459326",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 27: Sample two values from uniform distribution between 1 and 2 using `key_4` and assign it to `sample_uniform`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "7b58e260",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:40.929433Z",
     "iopub.status.busy": "2022-04-07T09:11:40.906414Z",
     "iopub.status.idle": "2022-04-07T09:11:41.176350Z",
     "shell.execute_reply": "2022-04-07T09:11:41.177136Z",
     "shell.execute_reply.started": "2022-04-07T06:11:54.310942Z"
    },
    "papermill": {
     "duration": 0.420091,
     "end_time": "2022-04-07T09:11:41.177378",
     "exception": false,
     "start_time": "2022-04-07T09:11:40.757287",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1.6274643, 1.1133162], dtype=float32)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_uniform = jax.random.uniform(key_4, shape=(2,), minval=1, maxval=2)\n",
    "sample_uniform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d81eaeb9",
   "metadata": {
    "papermill": {
     "duration": 0.145406,
     "end_time": "2022-04-07T09:11:41.468328",
     "exception": false,
     "start_time": "2022-04-07T09:11:41.322922",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 28: Sample three values from bernoulli distribution using `key_5` and assign it to `sample_bernoulli`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "07240595",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:41.791618Z",
     "iopub.status.busy": "2022-04-07T09:11:41.790662Z",
     "iopub.status.idle": "2022-04-07T09:11:42.092116Z",
     "shell.execute_reply": "2022-04-07T09:11:42.092979Z",
     "shell.execute_reply.started": "2022-04-07T06:11:54.437869Z"
    },
    "papermill": {
     "duration": 0.462377,
     "end_time": "2022-04-07T09:11:42.093208",
     "exception": false,
     "start_time": "2022-04-07T09:11:41.630831",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([False,  True,  True], dtype=bool)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_bernoulli = jax.random.bernoulli(key_5, shape=(3,))\n",
    "sample_bernoulli"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1349da3f",
   "metadata": {
    "papermill": {
     "duration": 0.150654,
     "end_time": "2022-04-07T09:11:42.390327",
     "exception": false,
     "start_time": "2022-04-07T09:11:42.239673",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 29: Sample a 2x3 matrix from poisson distribution with λ=100 using `key_6` and assign it to `sample_poisson`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "17667405",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:42.690846Z",
     "iopub.status.busy": "2022-04-07T09:11:42.689878Z",
     "iopub.status.idle": "2022-04-07T09:11:44.350660Z",
     "shell.execute_reply": "2022-04-07T09:11:44.351346Z",
     "shell.execute_reply.started": "2022-04-07T06:11:54.647155Z"
    },
    "papermill": {
     "duration": 1.813243,
     "end_time": "2022-04-07T09:11:44.351606",
     "exception": false,
     "start_time": "2022-04-07T09:11:42.538363",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 88,  82, 110],\n",
       "             [ 89,  85,  98]], dtype=int32)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_poisson = jax.random.poisson(key_6, shape=(2, 3), lam=100)\n",
    "sample_poisson"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32474d99",
   "metadata": {
    "papermill": {
     "duration": 0.149846,
     "end_time": "2022-04-07T09:11:44.647995",
     "exception": false,
     "start_time": "2022-04-07T09:11:44.498149",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 30: Sample a 2x3x4 array from normal distribution using `key_7` and assign it to `sample_normal`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "12e75c61",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:45.004979Z",
     "iopub.status.busy": "2022-04-07T09:11:44.963857Z",
     "iopub.status.idle": "2022-04-07T09:11:45.264099Z",
     "shell.execute_reply": "2022-04-07T09:11:45.264862Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.601132Z"
    },
    "papermill": {
     "duration": 0.467584,
     "end_time": "2022-04-07T09:11:45.265094",
     "exception": false,
     "start_time": "2022-04-07T09:11:44.797510",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[[ 0.25418088,  1.1962731 ,  1.3234351 ,  0.79711384],\n",
       "              [-1.8524722 , -0.28634202,  0.2251514 , -0.6195333 ],\n",
       "              [ 2.4013765 ,  0.07618266,  1.2277839 , -0.7562425 ]],\n",
       "\n",
       "             [[-0.45340484,  1.1029627 , -0.39860612, -1.1235143 ],\n",
       "              [-1.5689532 ,  0.4617323 , -0.5607138 , -1.7508575 ],\n",
       "              [ 0.50200105, -1.4972546 , -1.6995528 ,  0.5555226 ]]],            dtype=float32)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_normal = jax.random.normal(key_7, shape=(2, 3, 4))\n",
    "sample_normal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20eb429e",
   "metadata": {
    "papermill": {
     "duration": 0.14907,
     "end_time": "2022-04-07T09:11:45.560422",
     "exception": false,
     "start_time": "2022-04-07T09:11:45.411352",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 31: JIT-compile the `cube` function and assign it to `cube_jit`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "5a6e6919",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:45.862377Z",
     "iopub.status.busy": "2022-04-07T09:11:45.861366Z",
     "iopub.status.idle": "2022-04-07T09:11:45.865297Z",
     "shell.execute_reply": "2022-04-07T09:11:45.866056Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.918293Z"
    },
    "papermill": {
     "duration": 0.159522,
     "end_time": "2022-04-07T09:11:45.866297",
     "exception": false,
     "start_time": "2022-04-07T09:11:45.706775",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CompiledFunction at 0x7f8184303d70>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cube_jit = jax.jit(cube)\n",
    "cube_jit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3baf3ec",
   "metadata": {
    "papermill": {
     "duration": 0.148528,
     "end_time": "2022-04-07T09:11:46.169205",
     "exception": false,
     "start_time": "2022-04-07T09:11:46.020677",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 32: Display the execution time of `cube_jit` for first run (with overhead) with input=10.24**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "b7209c5d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:46.487897Z",
     "iopub.status.busy": "2022-04-07T09:11:46.486910Z",
     "iopub.status.idle": "2022-04-07T09:11:46.519414Z",
     "shell.execute_reply": "2022-04-07T09:11:46.518604Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.926521Z"
    },
    "papermill": {
     "duration": 0.197613,
     "end_time": "2022-04-07T09:11:46.519646",
     "exception": false,
     "start_time": "2022-04-07T09:11:46.322033",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 19.2 ms, sys: 11.1 ms, total: 30.3 ms\n",
      "Wall time: 26 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(1073.7418, dtype=float32)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "cube_jit(10.24)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d1e7ed",
   "metadata": {
    "papermill": {
     "duration": 0.153464,
     "end_time": "2022-04-07T09:11:46.827807",
     "exception": false,
     "start_time": "2022-04-07T09:11:46.674343",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 33: Display execution time of `cube_jit` for second run (without overhead) with input=10.24**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "49c8799b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:47.139958Z",
     "iopub.status.busy": "2022-04-07T09:11:47.138783Z",
     "iopub.status.idle": "2022-04-07T09:11:47.148423Z",
     "shell.execute_reply": "2022-04-07T09:11:47.149176Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.957502Z"
    },
    "papermill": {
     "duration": 0.169113,
     "end_time": "2022-04-07T09:11:47.149425",
     "exception": false,
     "start_time": "2022-04-07T09:11:46.980312",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.08 ms, sys: 2.02 ms, total: 4.1 ms\n",
      "Wall time: 2.63 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(1073.7418, dtype=float32)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "cube_jit(10.24)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "726a3588",
   "metadata": {
    "papermill": {
     "duration": 0.150367,
     "end_time": "2022-04-07T09:11:47.450467",
     "exception": false,
     "start_time": "2022-04-07T09:11:47.300100",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 34: Run `cube_jit` with input=10.24 and assign it to `cube_value`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "51a2aee4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:47.760431Z",
     "iopub.status.busy": "2022-04-07T09:11:47.758164Z",
     "iopub.status.idle": "2022-04-07T09:11:47.770841Z",
     "shell.execute_reply": "2022-04-07T09:11:47.771649Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.968435Z"
    },
    "papermill": {
     "duration": 0.17427,
     "end_time": "2022-04-07T09:11:47.771894",
     "exception": false,
     "start_time": "2022-04-07T09:11:47.597624",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(1073.7418, dtype=float32)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cube_value = cube_jit(10.24)\n",
    "cube_value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1d85e97",
   "metadata": {
    "papermill": {
     "duration": 0.154772,
     "end_time": "2022-04-07T09:11:48.075660",
     "exception": false,
     "start_time": "2022-04-07T09:11:47.920888",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 35: Run `cube_jit` with jit disabled and input=10.24 and assign it to `cube_value_nojit`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "c8ef4220",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:48.385579Z",
     "iopub.status.busy": "2022-04-07T09:11:48.384544Z",
     "iopub.status.idle": "2022-04-07T09:11:48.388085Z",
     "shell.execute_reply": "2022-04-07T09:11:48.388879Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.980573Z"
    },
    "papermill": {
     "duration": 0.164861,
     "end_time": "2022-04-07T09:11:48.389114",
     "exception": false,
     "start_time": "2022-04-07T09:11:48.224253",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1073.7418240000002"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with jax.disable_jit():\n",
    "    cube_value_nojit = cube_jit(10.24)\n",
    "\n",
    "cube_value_nojit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e72f259",
   "metadata": {
    "papermill": {
     "duration": 0.154968,
     "end_time": "2022-04-07T09:11:48.693654",
     "exception": false,
     "start_time": "2022-04-07T09:11:48.538686",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 36: Evaluate the shape of `cube_jit` with input=10.24 and assign it to `cube_shape`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "f38e5f3d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:49.006199Z",
     "iopub.status.busy": "2022-04-07T09:11:49.005205Z",
     "iopub.status.idle": "2022-04-07T09:11:49.009970Z",
     "shell.execute_reply": "2022-04-07T09:11:49.009081Z",
     "shell.execute_reply.started": "2022-04-07T06:11:57.993367Z"
    },
    "papermill": {
     "duration": 0.161627,
     "end_time": "2022-04-07T09:11:49.010165",
     "exception": false,
     "start_time": "2022-04-07T09:11:48.848538",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ShapeDtypeStruct(shape=(), dtype=float32)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cube_shape = jax.eval_shape(cube_jit, 10.24)\n",
    "cube_shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "510b8e9c",
   "metadata": {
    "papermill": {
     "duration": 0.148883,
     "end_time": "2022-04-07T09:11:49.306158",
     "exception": false,
     "start_time": "2022-04-07T09:11:49.157275",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 37: Create the jaxpr of `cube_jit` with input=10.24 and assign it to `cube_jaxpr`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "c8e166d5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:49.618287Z",
     "iopub.status.busy": "2022-04-07T09:11:49.617278Z",
     "iopub.status.idle": "2022-04-07T09:11:49.622609Z",
     "shell.execute_reply": "2022-04-07T09:11:49.621794Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.009397Z"
    },
    "papermill": {
     "duration": 0.16518,
     "end_time": "2022-04-07T09:11:49.622809",
     "exception": false,
     "start_time": "2022-04-07T09:11:49.457629",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ lambda  ; a.\n",
       "  let b = xla_call[ backend=None\n",
       "                    call_jaxpr={ lambda  ; a.\n",
       "                                 let b = integer_pow[ y=3 ] a\n",
       "                                 in (b,) }\n",
       "                    device=None\n",
       "                    donated_invars=(False,)\n",
       "                    inline=False\n",
       "                    name=cube ] a\n",
       "  in (b,) }"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cube_jaxpr = jax.make_jaxpr(cube_jit)(10.24)\n",
    "cube_jaxpr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7789311",
   "metadata": {
    "papermill": {
     "duration": 0.155919,
     "end_time": "2022-04-07T09:11:49.928144",
     "exception": false,
     "start_time": "2022-04-07T09:11:49.772225",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 38: Assign the XLA computation of `cube_jit` with input=10.24 to `cube_xla` and print it's XLA HLO text**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "483d4f72",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:50.262826Z",
     "iopub.status.busy": "2022-04-07T09:11:50.261819Z",
     "iopub.status.idle": "2022-04-07T09:11:50.266071Z",
     "shell.execute_reply": "2022-04-07T09:11:50.266883Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.027174Z"
    },
    "papermill": {
     "duration": 0.18545,
     "end_time": "2022-04-07T09:11:50.267132",
     "exception": false,
     "start_time": "2022-04-07T09:11:50.081682",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HloModule xla_computation_cube.12\n",
      "\n",
      "jit_cube__1.3 {\n",
      "  constant.5 = pred[] constant(false)\n",
      "  parameter.4 = f32[] parameter(0)\n",
      "  multiply.6 = f32[] multiply(parameter.4, parameter.4)\n",
      "  multiply.7 = f32[] multiply(parameter.4, multiply.6)\n",
      "  ROOT tuple.8 = (f32[]) tuple(multiply.7)\n",
      "}\n",
      "\n",
      "ENTRY xla_computation_cube.12 {\n",
      "  constant.2 = pred[] constant(false)\n",
      "  parameter.1 = f32[] parameter(0)\n",
      "  call.9 = (f32[]) call(parameter.1), to_apply=jit_cube__1.3\n",
      "  get-tuple-element.10 = f32[] get-tuple-element(call.9), index=0\n",
      "  ROOT tuple.11 = (f32[]) tuple(get-tuple-element.10)\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cube_xla = jax.xla_computation(cube_jit)(10.24)\n",
    "print(cube_xla.as_hlo_text())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271f9bff",
   "metadata": {
    "papermill": {
     "duration": 0.151608,
     "end_time": "2022-04-07T09:11:50.569435",
     "exception": false,
     "start_time": "2022-04-07T09:11:50.417827",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 39: Use the name `jaxton_cube_fn` internally for the `cube_jit` function and assign the named function to `cube_named_jit`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "57bae3c3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:50.884044Z",
     "iopub.status.busy": "2022-04-07T09:11:50.883032Z",
     "iopub.status.idle": "2022-04-07T09:11:50.886851Z",
     "shell.execute_reply": "2022-04-07T09:11:50.887607Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.045408Z"
    },
    "papermill": {
     "duration": 0.164467,
     "end_time": "2022-04-07T09:11:50.887837",
     "exception": false,
     "start_time": "2022-04-07T09:11:50.723370",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.cube(x)>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cube_named_jit = jax.named_call(cube_jit, name='jaxton_cube_fn')\n",
    "cube_named_jit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa4d32a",
   "metadata": {
    "papermill": {
     "duration": 0.15232,
     "end_time": "2022-04-07T09:11:51.194928",
     "exception": false,
     "start_time": "2022-04-07T09:11:51.042608",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 40: Assign the XLA computation of `cube_named_jit` with input=10.24 to `cube_named_xla` and print it's XLA HLO text**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9bd1fe17",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:51.515731Z",
     "iopub.status.busy": "2022-04-07T09:11:51.508123Z",
     "iopub.status.idle": "2022-04-07T09:11:51.520130Z",
     "shell.execute_reply": "2022-04-07T09:11:51.520991Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.059532Z"
    },
    "papermill": {
     "duration": 0.172553,
     "end_time": "2022-04-07T09:11:51.521248",
     "exception": false,
     "start_time": "2022-04-07T09:11:51.348695",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HloModule xla_computation_cube__1.18\n",
      "\n",
      "jit_cube__2.3 {\n",
      "  constant.5 = pred[] constant(false)\n",
      "  parameter.4 = f32[] parameter(0)\n",
      "  multiply.6 = f32[] multiply(parameter.4, parameter.4)\n",
      "  multiply.7 = f32[] multiply(parameter.4, multiply.6)\n",
      "  ROOT tuple.8 = (f32[]) tuple(multiply.7)\n",
      "}\n",
      "\n",
      "jaxton_cube_fn.9 {\n",
      "  constant.11 = pred[] constant(false)\n",
      "  parameter.10 = f32[] parameter(0)\n",
      "  call.12 = (f32[]) call(parameter.10), to_apply=jit_cube__2.3\n",
      "  get-tuple-element.13 = f32[] get-tuple-element(call.12), index=0\n",
      "  ROOT tuple.14 = (f32[]) tuple(get-tuple-element.13)\n",
      "}\n",
      "\n",
      "ENTRY xla_computation_cube__1.18 {\n",
      "  constant.2 = pred[] constant(false)\n",
      "  parameter.1 = f32[] parameter(0)\n",
      "  call.15 = (f32[]) call(parameter.1), to_apply=jaxton_cube_fn.9\n",
      "  get-tuple-element.16 = f32[] get-tuple-element(call.15), index=0\n",
      "  ROOT tuple.17 = (f32[]) tuple(get-tuple-element.16)\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cube_named_xla = jax.xla_computation(cube_named_jit)(10.24)\n",
    "print(cube_named_xla.as_hlo_text())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef301df7",
   "metadata": {
    "papermill": {
     "duration": 0.155826,
     "end_time": "2022-04-07T09:11:51.833657",
     "exception": false,
     "start_time": "2022-04-07T09:11:51.677831",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 41: Calculate the cumulative sum of `sample_data` using the associative scan operator and assign it to `data_cumsum`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "bc030004",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:52.154403Z",
     "iopub.status.busy": "2022-04-07T09:11:52.153367Z",
     "iopub.status.idle": "2022-04-07T09:11:52.571633Z",
     "shell.execute_reply": "2022-04-07T09:11:52.572398Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.077266Z"
    },
    "papermill": {
     "duration": 0.582374,
     "end_time": "2022-04-07T09:11:52.572649",
     "exception": false,
     "start_time": "2022-04-07T09:11:51.990275",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([10, 11, 35, 55, 70, 84], dtype=int32)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_cumsum = jax.lax.associative_scan(jnp.add, sample_data)\n",
    "data_cumsum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb4fe5ae",
   "metadata": {
    "papermill": {
     "duration": 0.161004,
     "end_time": "2022-04-07T09:11:52.911143",
     "exception": false,
     "start_time": "2022-04-07T09:11:52.750139",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 42: Calculate the cumulative sum of `sample_data` in reverse order using the associative scan operator and assign it to `data_cumsum_reverse`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "8e2ddca0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:53.232703Z",
     "iopub.status.busy": "2022-04-07T09:11:53.231781Z",
     "iopub.status.idle": "2022-04-07T09:11:53.294032Z",
     "shell.execute_reply": "2022-04-07T09:11:53.294831Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.329463Z"
    },
    "papermill": {
     "duration": 0.228166,
     "end_time": "2022-04-07T09:11:53.295114",
     "exception": false,
     "start_time": "2022-04-07T09:11:53.066948",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([84, 74, 73, 49, 29, 14], dtype=int32)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_cumsum_reverse = jax.lax.associative_scan(jnp.add, sample_data, reverse=True)\n",
    "data_cumsum_reverse"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9769d2a7",
   "metadata": {
    "papermill": {
     "duration": 0.161764,
     "end_time": "2022-04-07T09:11:53.613869",
     "exception": false,
     "start_time": "2022-04-07T09:11:53.452105",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 43: Create a JIT-compiled lambda function that outputs `square` of input if it is even and `cube` of input if it is odd using the cond operator and assign it to `parity_ifelse`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "26d85456",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:53.933892Z",
     "iopub.status.busy": "2022-04-07T09:11:53.932914Z",
     "iopub.status.idle": "2022-04-07T09:11:53.936483Z",
     "shell.execute_reply": "2022-04-07T09:11:53.937206Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.353657Z"
    },
    "papermill": {
     "duration": 0.169787,
     "end_time": "2022-04-07T09:11:53.937433",
     "exception": false,
     "start_time": "2022-04-07T09:11:53.767646",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CompiledFunction at 0x7f8184303c80>"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_ifelse = jax.jit(lambda x: jax.lax.cond(jnp.remainder(x, 2) == 0, square, cube, x))\n",
    "parity_ifelse"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "509dc218",
   "metadata": {
    "papermill": {
     "duration": 0.157647,
     "end_time": "2022-04-07T09:11:54.250182",
     "exception": false,
     "start_time": "2022-04-07T09:11:54.092535",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 44: Run `parity_ifelse` with the first element of `data_cumsum` and assign it to `parity_1`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "6589b821",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:54.571246Z",
     "iopub.status.busy": "2022-04-07T09:11:54.570032Z",
     "iopub.status.idle": "2022-04-07T09:11:54.629404Z",
     "shell.execute_reply": "2022-04-07T09:11:54.628623Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.362631Z"
    },
    "papermill": {
     "duration": 0.22604,
     "end_time": "2022-04-07T09:11:54.629620",
     "exception": false,
     "start_time": "2022-04-07T09:11:54.403580",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(100, dtype=int32)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_1 = parity_ifelse(data_cumsum[0])\n",
    "parity_1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84ae38d5",
   "metadata": {
    "papermill": {
     "duration": 0.15497,
     "end_time": "2022-04-07T09:11:54.946967",
     "exception": false,
     "start_time": "2022-04-07T09:11:54.791997",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 45: Run `parity_ifelse` with the second element of `data_cumsum` and assign it to `parity_2`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "eaceeeea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:55.269155Z",
     "iopub.status.busy": "2022-04-07T09:11:55.267870Z",
     "iopub.status.idle": "2022-04-07T09:11:55.283751Z",
     "shell.execute_reply": "2022-04-07T09:11:55.284499Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.421739Z"
    },
    "papermill": {
     "duration": 0.178933,
     "end_time": "2022-04-07T09:11:55.284753",
     "exception": false,
     "start_time": "2022-04-07T09:11:55.105820",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(1331, dtype=int32)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_2 = parity_ifelse(data_cumsum[1])\n",
    "parity_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "227479fe",
   "metadata": {
    "papermill": {
     "duration": 0.167085,
     "end_time": "2022-04-07T09:11:55.614641",
     "exception": false,
     "start_time": "2022-04-07T09:11:55.447556",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 46: Create a JIT-compiled lambda function that outputs `square` of input if it is even and `cube` of input if it is odd using the switch operator and assign it to `parity_switch`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "5cf46ac6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:55.940601Z",
     "iopub.status.busy": "2022-04-07T09:11:55.939617Z",
     "iopub.status.idle": "2022-04-07T09:11:55.944278Z",
     "shell.execute_reply": "2022-04-07T09:11:55.943390Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.432623Z"
    },
    "papermill": {
     "duration": 0.169882,
     "end_time": "2022-04-07T09:11:55.944481",
     "exception": false,
     "start_time": "2022-04-07T09:11:55.774599",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CompiledFunction at 0x7f8184303aa0>"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_switch = jax.jit(lambda x: jax.lax.switch(jnp.remainder(x, 2), [square, cube], x))\n",
    "parity_switch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a830233",
   "metadata": {
    "papermill": {
     "duration": 0.160608,
     "end_time": "2022-04-07T09:11:56.264899",
     "exception": false,
     "start_time": "2022-04-07T09:11:56.104291",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 47: Run `parity_switch` with the fourth element of `data_cumsum` and assign it to `parity_4`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "2804c8d2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:56.603530Z",
     "iopub.status.busy": "2022-04-07T09:11:56.602487Z",
     "iopub.status.idle": "2022-04-07T09:11:56.662382Z",
     "shell.execute_reply": "2022-04-07T09:11:56.663127Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.447789Z"
    },
    "papermill": {
     "duration": 0.231517,
     "end_time": "2022-04-07T09:11:56.663371",
     "exception": false,
     "start_time": "2022-04-07T09:11:56.431854",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(166375, dtype=int32)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_4 = parity_switch(data_cumsum[3])\n",
    "parity_4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "696d2b31",
   "metadata": {
    "papermill": {
     "duration": 0.161685,
     "end_time": "2022-04-07T09:11:56.983945",
     "exception": false,
     "start_time": "2022-04-07T09:11:56.822260",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 48: Run `parity_switch` with the fifth element of `data_cumsum` and assign it to `parity_5`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "5e8b73c0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:57.309718Z",
     "iopub.status.busy": "2022-04-07T09:11:57.307870Z",
     "iopub.status.idle": "2022-04-07T09:11:57.324590Z",
     "shell.execute_reply": "2022-04-07T09:11:57.325328Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.492597Z"
    },
    "papermill": {
     "duration": 0.183594,
     "end_time": "2022-04-07T09:11:57.325573",
     "exception": false,
     "start_time": "2022-04-07T09:11:57.141979",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(4900, dtype=int32)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parity_5 = parity_switch(data_cumsum[4])\n",
    "parity_5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b60ac986",
   "metadata": {
    "papermill": {
     "duration": 0.159803,
     "end_time": "2022-04-07T09:11:57.644307",
     "exception": false,
     "start_time": "2022-04-07T09:11:57.484504",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 49: Calculate the sum of the first four elements of `data_cumsum` using the for operator and assign it to `sum_four`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "5ff8b2ab",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:57.983944Z",
     "iopub.status.busy": "2022-04-07T09:11:57.964272Z",
     "iopub.status.idle": "2022-04-07T09:11:58.028688Z",
     "shell.execute_reply": "2022-04-07T09:11:58.029773Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.502732Z"
    },
    "papermill": {
     "duration": 0.230934,
     "end_time": "2022-04-07T09:11:58.030088",
     "exception": false,
     "start_time": "2022-04-07T09:11:57.799154",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(111, dtype=int32)"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum_four = jax.lax.fori_loop(0, 4, lambda i, x: x+data_cumsum[i], 0)\n",
    "sum_four"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc373b01",
   "metadata": {
    "papermill": {
     "duration": 0.158089,
     "end_time": "2022-04-07T09:11:58.350205",
     "exception": false,
     "start_time": "2022-04-07T09:11:58.192116",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 50: Keep subtracting 25 from `sum_four` until the result is negative using the while operator and assign it to `subtract_until_negative`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "0f4e784c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:58.690622Z",
     "iopub.status.busy": "2022-04-07T09:11:58.689675Z",
     "iopub.status.idle": "2022-04-07T09:11:58.725522Z",
     "shell.execute_reply": "2022-04-07T09:11:58.724686Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.561559Z"
    },
    "papermill": {
     "duration": 0.208402,
     "end_time": "2022-04-07T09:11:58.725723",
     "exception": false,
     "start_time": "2022-04-07T09:11:58.517321",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(-14, dtype=int32)"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subtract_until_negative = jax.lax.while_loop(lambda x: x>0, lambda x: x-25, sum_four)\n",
    "subtract_until_negative"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29704790",
   "metadata": {
    "papermill": {
     "duration": 0.159313,
     "end_time": "2022-04-07T09:11:59.048631",
     "exception": false,
     "start_time": "2022-04-07T09:11:58.889318",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 51: JIT-compile the derivative of `cube` and assign it to `derivative_cube`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "8bcff12a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:11:59.377160Z",
     "iopub.status.busy": "2022-04-07T09:11:59.376124Z",
     "iopub.status.idle": "2022-04-07T09:11:59.379978Z",
     "shell.execute_reply": "2022-04-07T09:11:59.380716Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.598340Z"
    },
    "papermill": {
     "duration": 0.170411,
     "end_time": "2022-04-07T09:11:59.380953",
     "exception": false,
     "start_time": "2022-04-07T09:11:59.210542",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CompiledFunction at 0x7f81841edb90>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_cube = jax.jit(jax.grad(cube))\n",
    "derivative_cube"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8040cf46",
   "metadata": {
    "papermill": {
     "duration": 0.164007,
     "end_time": "2022-04-07T09:11:59.700897",
     "exception": false,
     "start_time": "2022-04-07T09:11:59.536890",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 52: Run `derivative_cube` with value=7**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "fa3cb4c3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:00.037235Z",
     "iopub.status.busy": "2022-04-07T09:12:00.036240Z",
     "iopub.status.idle": "2022-04-07T09:12:00.065128Z",
     "shell.execute_reply": "2022-04-07T09:12:00.066178Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.605259Z"
    },
    "papermill": {
     "duration": 0.20407,
     "end_time": "2022-04-07T09:12:00.066434",
     "exception": false,
     "start_time": "2022-04-07T09:11:59.862364",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(147., dtype=float32)"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_cube(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e31e524",
   "metadata": {
    "papermill": {
     "duration": 0.157865,
     "end_time": "2022-04-07T09:12:00.389069",
     "exception": false,
     "start_time": "2022-04-07T09:12:00.231204",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 53: JIT-compile the value and derivative of `cube` together, assign it to `value_and_derivative_cube` and run it with value=7**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "cbcb6c4b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:00.727208Z",
     "iopub.status.busy": "2022-04-07T09:12:00.726253Z",
     "iopub.status.idle": "2022-04-07T09:12:00.756652Z",
     "shell.execute_reply": "2022-04-07T09:12:00.757417Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.640352Z"
    },
    "papermill": {
     "duration": 0.207446,
     "end_time": "2022-04-07T09:12:00.757671",
     "exception": false,
     "start_time": "2022-04-07T09:12:00.550225",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(343., dtype=float32), DeviceArray(147., dtype=float32))"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "value_and_derivative_cube = jax.jit(jax.value_and_grad(cube))\n",
    "value_and_derivative_cube(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26fa3c6",
   "metadata": {
    "papermill": {
     "duration": 0.16308,
     "end_time": "2022-04-07T09:12:01.086515",
     "exception": false,
     "start_time": "2022-04-07T09:12:00.923435",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 54: JIT-compile the second order derivative of `cube`, assign it to `derivative_cube_2` and run it with value=7**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "ad0829b3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:01.435732Z",
     "iopub.status.busy": "2022-04-07T09:12:01.417187Z",
     "iopub.status.idle": "2022-04-07T09:12:01.467902Z",
     "shell.execute_reply": "2022-04-07T09:12:01.467108Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.669500Z"
    },
    "papermill": {
     "duration": 0.215853,
     "end_time": "2022-04-07T09:12:01.468108",
     "exception": false,
     "start_time": "2022-04-07T09:12:01.252255",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(42., dtype=float32)"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_cube_2 = jax.jit(jax.grad(jax.grad(cube)))\n",
    "derivative_cube_2(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cf63452",
   "metadata": {
    "papermill": {
     "duration": 0.164437,
     "end_time": "2022-04-07T09:12:01.794672",
     "exception": false,
     "start_time": "2022-04-07T09:12:01.630235",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 55: JIT-compile the hessian of `cube`, assign it to `hessian_cube` and run it with value=7**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "bd727f5e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:02.130651Z",
     "iopub.status.busy": "2022-04-07T09:12:02.129340Z",
     "iopub.status.idle": "2022-04-07T09:12:02.180775Z",
     "shell.execute_reply": "2022-04-07T09:12:02.179922Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.704309Z"
    },
    "papermill": {
     "duration": 0.223161,
     "end_time": "2022-04-07T09:12:02.180992",
     "exception": false,
     "start_time": "2022-04-07T09:12:01.957831",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(42., dtype=float32)"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hessian_cube = jax.jit(jax.hessian(cube))\n",
    "hessian_cube(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3eca93a",
   "metadata": {
    "papermill": {
     "duration": 0.164783,
     "end_time": "2022-04-07T09:12:02.514811",
     "exception": false,
     "start_time": "2022-04-07T09:12:02.350028",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 56: JIT-compile `areas`, assign it to `jit_areas` and run it with value=9**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "5e385e99",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:02.850128Z",
     "iopub.status.busy": "2022-04-07T09:12:02.849095Z",
     "iopub.status.idle": "2022-04-07T09:12:02.891491Z",
     "shell.execute_reply": "2022-04-07T09:12:02.892302Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.756807Z"
    },
    "papermill": {
     "duration": 0.219437,
     "end_time": "2022-04-07T09:12:02.892557",
     "exception": false,
     "start_time": "2022-04-07T09:12:02.673120",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[DeviceArray(35.074028, dtype=float32),\n",
       " DeviceArray(81, dtype=int32),\n",
       " DeviceArray(254.46901, dtype=float32)]"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jit_areas = jax.jit(areas)\n",
    "jit_areas(9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12870309",
   "metadata": {
    "papermill": {
     "duration": 0.164692,
     "end_time": "2022-04-07T09:12:03.222323",
     "exception": false,
     "start_time": "2022-04-07T09:12:03.057631",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 57: Compute the Jacobian of `areas` using forward-mode automatic differentiation, assign it to `jacfwd_areas` and run it with value=9**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "5fb692bd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:03.558518Z",
     "iopub.status.busy": "2022-04-07T09:12:03.557343Z",
     "iopub.status.idle": "2022-04-07T09:12:03.759817Z",
     "shell.execute_reply": "2022-04-07T09:12:03.760864Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.789504Z"
    },
    "papermill": {
     "duration": 0.373186,
     "end_time": "2022-04-07T09:12:03.761210",
     "exception": false,
     "start_time": "2022-04-07T09:12:03.388024",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[DeviceArray(7.7942286, dtype=float32),\n",
       " DeviceArray(18., dtype=float32),\n",
       " DeviceArray(56.548668, dtype=float32)]"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacfwd_areas = jax.jacfwd(areas)\n",
    "jacfwd_areas(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d6c0ff5",
   "metadata": {
    "papermill": {
     "duration": 0.168248,
     "end_time": "2022-04-07T09:12:04.095805",
     "exception": false,
     "start_time": "2022-04-07T09:12:03.927557",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 58: Compute the Jacobian of `areas` using reverse-mode automatic differentiation, assign it to `jacrev_areas` and run it with value=9**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "b9cd8835",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:04.430267Z",
     "iopub.status.busy": "2022-04-07T09:12:04.429252Z",
     "iopub.status.idle": "2022-04-07T09:12:04.676007Z",
     "shell.execute_reply": "2022-04-07T09:12:04.675197Z",
     "shell.execute_reply.started": "2022-04-07T06:11:58.980002Z"
    },
    "papermill": {
     "duration": 0.415866,
     "end_time": "2022-04-07T09:12:04.676222",
     "exception": false,
     "start_time": "2022-04-07T09:12:04.260356",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[DeviceArray(7.7942286, dtype=float32),\n",
       " DeviceArray(18., dtype=float32),\n",
       " DeviceArray(56.548668, dtype=float32)]"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacrev_areas = jax.jacrev(areas)\n",
    "jacrev_areas(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c530ad7",
   "metadata": {
    "papermill": {
     "duration": 0.174219,
     "end_time": "2022-04-07T09:12:05.016389",
     "exception": false,
     "start_time": "2022-04-07T09:12:04.842170",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 59: Compute the Jacobian-vector product of `cube` at value=7 with vector=9 and assign it to `jvp_cube`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "0ca73bb2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:05.357719Z",
     "iopub.status.busy": "2022-04-07T09:12:05.356658Z",
     "iopub.status.idle": "2022-04-07T09:12:05.390981Z",
     "shell.execute_reply": "2022-04-07T09:12:05.391710Z",
     "shell.execute_reply.started": "2022-04-07T06:18:41.707168Z"
    },
    "papermill": {
     "duration": 0.207371,
     "end_time": "2022-04-07T09:12:05.391957",
     "exception": false,
     "start_time": "2022-04-07T09:12:05.184586",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(343., dtype=float32), DeviceArray(1323., dtype=float32))"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jvp_cube = jax.jvp(cube, (7.0,), (9.0,))\n",
    "jvp_cube"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda2afa3",
   "metadata": {
    "papermill": {
     "duration": 0.167241,
     "end_time": "2022-04-07T09:12:05.740273",
     "exception": false,
     "start_time": "2022-04-07T09:12:05.573032",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 60: Compute the linear approximation of `areas` with value=5, assign it to `areas_linear` and run it with value=9**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "94de2127",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T09:12:06.058185Z",
     "iopub.status.busy": "2022-04-07T09:12:06.054236Z",
     "iopub.status.idle": "2022-04-07T09:12:06.178446Z",
     "shell.execute_reply": "2022-04-07T09:12:06.179247Z",
     "shell.execute_reply.started": "2022-04-07T06:11:59.241592Z"
    },
    "papermill": {
     "duration": 0.301753,
     "end_time": "2022-04-07T09:12:06.179495",
     "exception": false,
     "start_time": "2022-04-07T09:12:05.877742",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[DeviceArray(38.97114, dtype=float32),\n",
       " DeviceArray(90., dtype=float32),\n",
       " DeviceArray(282.74335, dtype=float32)]"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_, areas_linear = jax.linearize(areas, 5.0)\n",
    "areas_linear(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec8c197",
   "metadata": {
    "papermill": {
     "duration": 0.169417,
     "end_time": "2022-04-07T09:12:06.524853",
     "exception": false,
     "start_time": "2022-04-07T09:12:06.355436",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center>\n",
    "    All solutions of <b>JaxTon</b>: <i>💯 JAX exercises</i>\n",
    "    <br>\n",
    "    You can find individual sets of the exercises and solutions on <a href=\"https://github.com/vopani/jaxton#exercises-\">GitHub</a>\n",
    "</center>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.7.10"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 191.226993,
   "end_time": "2022-04-07T09:12:10.088352",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2022-04-07T09:08:58.861359",
   "version": "2.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}