notebooks/06_automatic_differentiation_exercises.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","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>","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":[]}},{"cell_type":"markdown","source":"<center>\n    This is Set 6: Automatic  Differentiation (Exercises 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":{"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":[]}},{"cell_type":"markdown","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.","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":[]}},{"cell_type":"code","source":"!python3 -m pip install jax","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import jax\nimport jax.numpy as jnp\nimport math\nimport os\nimport requests\n\ntry:\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']\nexcept:\n    pass\n\njax.devices()","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## sample cube function\ndef cube(x):\n    return x**3\n\ncube(2)","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## sample areas function\ndef areas(x):\n    return [math.sqrt(3)*x**2/4, x**2, math.pi*x**2]\n\nareas(2)","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 51: JIT-compile the derivative of `cube` and assign it to `derivative_cube`**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 52: Run `derivative_cube` with value=7**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","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**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 54: JIT-compile the second order derivative of `cube`, assign it to `derivative_cube_2` and run it with value=7**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 55: JIT-compile the hessian of `cube`, assign it to `hessian_cube` and run it with value=7**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 56: JIT-compile `areas`, assign it to `jit_areas` and run it with value=9**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 57: Compute the Jacobian of `areas` using forward-mode automatic differentiation, assign it to `jacfwd_areas` and run it with value=9**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 58: Compute the Jacobian of `areas` using reverse-mode automatic differentiation, assign it to `jacrev_areas` and run it with value=9**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 59: Compute the Jacobian-vector product of `cube` at value=7 with vector=9 and assign it to `jvp_cube`**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 60: Compute the linear approximation of `areas` with value=5, assign it to `areas_linear` and run it with value=9**","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":[]}},{"cell_type":"code","source":"","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":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<center>\n    This completes Set 6: Automatic Differentiation (Exercises 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":{"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":[]}}]}