notebooks/06_automatic_differentiation_solutions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3a18d19c",
   "metadata": {
    "papermill": {
     "duration": 0.019546,
     "end_time": "2022-04-07T12:16:47.515898",
     "exception": false,
     "start_time": "2022-04-07T12:16:47.496352",
     "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": "ad794fd9",
   "metadata": {
    "papermill": {
     "duration": 0.018846,
     "end_time": "2022-04-07T12:16:47.553979",
     "exception": false,
     "start_time": "2022-04-07T12:16:47.535133",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center>\n",
    "    This is Set 6: Automatic  Differentiation (Solutions 51-60) of <b>JaxTon</b>: <i>💯 JAX exercises</i>\n",
    "    <br>\n",
    "    You can find all the exercises and solutions on <a href=\"https://github.com/vopani/jaxton#exercises-\">GitHub</a>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "708e9280",
   "metadata": {
    "papermill": {
     "duration": 0.018161,
     "end_time": "2022-04-07T12:16:47.591089",
     "exception": false,
     "start_time": "2022-04-07T12:16:47.572928",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Prerequisites**\n",
    "\n",
    "* The configuration of jax should be set as shown in the code snippet below in order to use TPUs.\n",
    "* Sample functions `cube` and `areas` will be used for the exercises."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "38884779",
   "metadata": {
    "_kg_hide-input": true,
    "_kg_hide-output": true,
    "execution": {
     "iopub.execute_input": "2022-04-07T12:16:47.639881Z",
     "iopub.status.busy": "2022-04-07T12:16:47.639189Z",
     "iopub.status.idle": "2022-04-07T12:16:56.717020Z",
     "shell.execute_reply": "2022-04-07T12:16:56.716025Z",
     "shell.execute_reply.started": "2022-04-07T06:25:25.564042Z"
    },
    "papermill": {
     "duration": 9.107442,
     "end_time": "2022-04-07T12:16:56.717195",
     "exception": false,
     "start_time": "2022-04-07T12:16:47.609753",
     "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: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.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: 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": "code",
   "execution_count": 2,
   "id": "75f21b26",
   "metadata": {
    "_kg_hide-input": true,
    "execution": {
     "iopub.execute_input": "2022-04-07T12:16:56.764367Z",
     "iopub.status.busy": "2022-04-07T12:16:56.763709Z",
     "iopub.status.idle": "2022-04-07T12:17:45.842774Z",
     "shell.execute_reply": "2022-04-07T12:17:45.843548Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.319365Z"
    },
    "papermill": {
     "duration": 49.106015,
     "end_time": "2022-04-07T12:17:45.843807",
     "exception": false,
     "start_time": "2022-04-07T12:16:56.737792",
     "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 jax\n",
    "import jax.numpy as jnp\n",
    "import math\n",
    "import os\n",
    "import requests\n",
    "\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": "503d9f85",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:17:45.915377Z",
     "iopub.status.busy": "2022-04-07T12:17:45.914080Z",
     "iopub.status.idle": "2022-04-07T12:17:45.923112Z",
     "shell.execute_reply": "2022-04-07T12:17:45.923911Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.330799Z"
    },
    "papermill": {
     "duration": 0.04921,
     "end_time": "2022-04-07T12:17:45.924151",
     "exception": false,
     "start_time": "2022-04-07T12:17:45.874941",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "execution_count": 3,
     "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": 4,
   "id": "97f1d519",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:17:45.996403Z",
     "iopub.status.busy": "2022-04-07T12:17:45.995344Z",
     "iopub.status.idle": "2022-04-07T12:17:45.999411Z",
     "shell.execute_reply": "2022-04-07T12:17:46.000098Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.344617Z"
    },
    "papermill": {
     "duration": 0.044958,
     "end_time": "2022-04-07T12:17:46.000348",
     "exception": false,
     "start_time": "2022-04-07T12:17:45.955390",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.7320508075688772, 4, 12.566370614359172]"
      ]
     },
     "execution_count": 4,
     "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": "8a61392c",
   "metadata": {
    "papermill": {
     "duration": 0.031167,
     "end_time": "2022-04-07T12:17:46.064507",
     "exception": false,
     "start_time": "2022-04-07T12:17:46.033340",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 51: JIT-compile the derivative of `cube` and assign it to `derivative_cube`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4c1c79fa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:17:46.136176Z",
     "iopub.status.busy": "2022-04-07T12:17:46.135129Z",
     "iopub.status.idle": "2022-04-07T12:17:46.139793Z",
     "shell.execute_reply": "2022-04-07T12:17:46.140540Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.357854Z"
    },
    "papermill": {
     "duration": 0.043939,
     "end_time": "2022-04-07T12:17:46.140771",
     "exception": false,
     "start_time": "2022-04-07T12:17:46.096832",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CompiledFunction at 0x7f416e57d8c0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_cube = jax.jit(jax.grad(cube))\n",
    "derivative_cube"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "894489b5",
   "metadata": {
    "papermill": {
     "duration": 0.034101,
     "end_time": "2022-04-07T12:17:46.208577",
     "exception": false,
     "start_time": "2022-04-07T12:17:46.174476",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**Exercise 52: Run `derivative_cube` with value=7**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e6ab154e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:17:46.306345Z",
     "iopub.status.busy": "2022-04-07T12:17:46.305339Z",
     "iopub.status.idle": "2022-04-07T12:18:48.063559Z",
     "shell.execute_reply": "2022-04-07T12:18:48.064324Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.370543Z"
    },
    "papermill": {
     "duration": 61.822266,
     "end_time": "2022-04-07T12:18:48.064648",
     "exception": false,
     "start_time": "2022-04-07T12:17:46.242382",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(147., dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_cube(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf9178af",
   "metadata": {
    "papermill": {
     "duration": 0.032159,
     "end_time": "2022-04-07T12:18:48.129679",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.097520",
     "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": 7,
   "id": "676ee777",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:48.207980Z",
     "iopub.status.busy": "2022-04-07T12:18:48.206940Z",
     "iopub.status.idle": "2022-04-07T12:18:48.282544Z",
     "shell.execute_reply": "2022-04-07T12:18:48.283808Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.745862Z"
    },
    "papermill": {
     "duration": 0.12182,
     "end_time": "2022-04-07T12:18:48.284057",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.162237",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(343., dtype=float32), DeviceArray(147., dtype=float32))"
      ]
     },
     "execution_count": 7,
     "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": "7f13008f",
   "metadata": {
    "papermill": {
     "duration": 0.032762,
     "end_time": "2022-04-07T12:18:48.350140",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.317378",
     "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": 8,
   "id": "45eba0bc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:48.444214Z",
     "iopub.status.busy": "2022-04-07T12:18:48.443258Z",
     "iopub.status.idle": "2022-04-07T12:18:48.481101Z",
     "shell.execute_reply": "2022-04-07T12:18:48.481945Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.777259Z"
    },
    "papermill": {
     "duration": 0.098391,
     "end_time": "2022-04-07T12:18:48.482325",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.383934",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(42., dtype=float32)"
      ]
     },
     "execution_count": 8,
     "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": "a594f80b",
   "metadata": {
    "papermill": {
     "duration": 0.034414,
     "end_time": "2022-04-07T12:18:48.552895",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.518481",
     "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": 9,
   "id": "c9ce344a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:48.644336Z",
     "iopub.status.busy": "2022-04-07T12:18:48.643095Z",
     "iopub.status.idle": "2022-04-07T12:18:48.682105Z",
     "shell.execute_reply": "2022-04-07T12:18:48.681114Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.815636Z"
    },
    "papermill": {
     "duration": 0.095631,
     "end_time": "2022-04-07T12:18:48.682324",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.586693",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(42., dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hessian_cube = jax.jit(jax.hessian(cube))\n",
    "hessian_cube(7.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ed7df8",
   "metadata": {
    "papermill": {
     "duration": 0.03569,
     "end_time": "2022-04-07T12:18:48.754701",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.719011",
     "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": 10,
   "id": "21ddbb8e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:48.841749Z",
     "iopub.status.busy": "2022-04-07T12:18:48.839259Z",
     "iopub.status.idle": "2022-04-07T12:18:48.873231Z",
     "shell.execute_reply": "2022-04-07T12:18:48.874008Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.867531Z"
    },
    "papermill": {
     "duration": 0.083448,
     "end_time": "2022-04-07T12:18:48.874244",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.790796",
     "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": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jit_areas = jax.jit(areas)\n",
    "jit_areas(9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "598a1729",
   "metadata": {
    "papermill": {
     "duration": 0.034341,
     "end_time": "2022-04-07T12:18:48.944932",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.910591",
     "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": 11,
   "id": "0d6eb211",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:49.026250Z",
     "iopub.status.busy": "2022-04-07T12:18:49.025317Z",
     "iopub.status.idle": "2022-04-07T12:18:49.218262Z",
     "shell.execute_reply": "2022-04-07T12:18:49.219501Z",
     "shell.execute_reply.started": "2022-04-07T06:25:34.905130Z"
    },
    "papermill": {
     "duration": 0.236784,
     "end_time": "2022-04-07T12:18:49.219774",
     "exception": false,
     "start_time": "2022-04-07T12:18:48.982990",
     "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": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacfwd_areas = jax.jacfwd(areas)\n",
    "jacfwd_areas(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce558699",
   "metadata": {
    "papermill": {
     "duration": 0.035467,
     "end_time": "2022-04-07T12:18:49.291957",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.256490",
     "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": 12,
   "id": "d7f75782",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:49.372867Z",
     "iopub.status.busy": "2022-04-07T12:18:49.371883Z",
     "iopub.status.idle": "2022-04-07T12:18:49.597907Z",
     "shell.execute_reply": "2022-04-07T12:18:49.598407Z",
     "shell.execute_reply.started": "2022-04-07T06:25:35.163099Z"
    },
    "papermill": {
     "duration": 0.270545,
     "end_time": "2022-04-07T12:18:49.598614",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.328069",
     "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": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacrev_areas = jax.jacrev(areas)\n",
    "jacrev_areas(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18db5d72",
   "metadata": {
    "papermill": {
     "duration": 0.036105,
     "end_time": "2022-04-07T12:18:49.670355",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.634250",
     "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": 13,
   "id": "fd51d6e2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:49.752480Z",
     "iopub.status.busy": "2022-04-07T12:18:49.751248Z",
     "iopub.status.idle": "2022-04-07T12:18:49.785271Z",
     "shell.execute_reply": "2022-04-07T12:18:49.786038Z",
     "shell.execute_reply.started": "2022-04-07T06:25:35.497031Z"
    },
    "papermill": {
     "duration": 0.080123,
     "end_time": "2022-04-07T12:18:49.786293",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.706170",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(343., dtype=float32), DeviceArray(1323., dtype=float32))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jvp_cube = jax.jvp(cube, (7.0,), (9.0,))\n",
    "jvp_cube"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82656bd",
   "metadata": {
    "papermill": {
     "duration": 0.039154,
     "end_time": "2022-04-07T12:18:49.870902",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.831748",
     "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": 14,
   "id": "2f9bbaa8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-04-07T12:18:49.954714Z",
     "iopub.status.busy": "2022-04-07T12:18:49.953453Z",
     "iopub.status.idle": "2022-04-07T12:18:50.063937Z",
     "shell.execute_reply": "2022-04-07T12:18:50.064630Z",
     "shell.execute_reply.started": "2022-04-07T06:25:35.523982Z"
    },
    "papermill": {
     "duration": 0.155407,
     "end_time": "2022-04-07T12:18:50.064869",
     "exception": false,
     "start_time": "2022-04-07T12:18:49.909462",
     "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": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_, areas_linear = jax.linearize(areas, 5.0)\n",
    "areas_linear(9.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6559fcc",
   "metadata": {
    "papermill": {
     "duration": 0.037608,
     "end_time": "2022-04-07T12:18:50.139612",
     "exception": false,
     "start_time": "2022-04-07T12:18:50.102004",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center>\n",
    "    This completes Set 6: Automatic Differentiation (Solutions 51-60) of <b>JaxTon</b>: <i>💯 JAX exercises</i>\n",
    "    <br>\n",
    "    You can find all 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": 133.719878,
   "end_time": "2022-04-07T12:18:53.257261",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2022-04-07T12:16:39.537383",
   "version": "2.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}