notebooks/01_jax_introduction_solutions.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"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.01716,"end_time":"2022-01-09T11:47:46.089693","exception":false,"start_time":"2022-01-09T11:47:46.072533","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"<center>\n    This is Set 1: JAX Introduction (Solutions 1-10) 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.015686,"end_time":"2022-01-09T11:47:46.12191","exception":false,"start_time":"2022-01-09T11:47:46.106224","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"**Prerequisites**\n\n* The configuration of jax should be updated as shown in the code snippet below in order to use TPUs.","metadata":{"papermill":{"duration":0.015589,"end_time":"2022-01-09T11:47:46.15344","exception":false,"start_time":"2022-01-09T11:47:46.137851","status":"completed"},"tags":[]}},{"cell_type":"code","source":"## install jax\n!python3 -m pip install jax","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-02-04T05:41:22.244841Z","iopub.execute_input":"2022-02-04T05:41:22.245537Z","iopub.status.idle":"2022-02-04T05:41:31.997684Z","shell.execute_reply.started":"2022-02-04T05:41:22.245425Z","shell.execute_reply":"2022-02-04T05:41:31.996427Z"},"trusted":true},"execution_count":1,"outputs":[{"name":"stdout","text":"Requirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (0.2.19)\nRequirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax) (3.3.0)\nRequirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\nRequirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\nRequirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from absl-py->jax) (1.15.0)\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\n","output_type":"stream"}]},{"cell_type":"code","source":"## import packages\nimport jax\nimport os\nimport requests\n\n## setup JAX to use TPUs if available\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":false,"papermill":{"duration":1.693423,"end_time":"2022-01-09T11:47:47.862757","exception":false,"start_time":"2022-01-09T11:47:46.169334","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:41:32.000225Z","iopub.execute_input":"2022-02-04T05:41:32.000545Z","iopub.status.idle":"2022-02-04T05:42:11.077482Z","shell.execute_reply.started":"2022-02-04T05:41:32.000506Z","shell.execute_reply":"2022-02-04T05:42:11.076034Z"},"trusted":true},"execution_count":2,"outputs":[{"execution_count":2,"output_type":"execute_result","data":{"text/plain":"[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 1: Install the `jax` package**","metadata":{"papermill":{"duration":0.015931,"end_time":"2022-01-09T11:47:47.894827","exception":false,"start_time":"2022-01-09T11:47:47.878896","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!python3 -m pip install jax","metadata":{"papermill":{"duration":9.887596,"end_time":"2022-01-09T11:47:57.799261","exception":false,"start_time":"2022-01-09T11:47:47.911665","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:11.079301Z","iopub.execute_input":"2022-02-04T05:42:11.079807Z","iopub.status.idle":"2022-02-04T05:42:23.055526Z","shell.execute_reply.started":"2022-02-04T05:42:11.079745Z","shell.execute_reply":"2022-02-04T05:42:23.053980Z"},"trusted":true},"execution_count":3,"outputs":[{"name":"stdout","text":"Requirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (0.2.19)\nRequirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\nRequirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\nRequirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax) (3.3.0)\nRequirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from absl-py->jax) (1.15.0)\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\n","output_type":"stream"}]},{"cell_type":"markdown","source":"**Exercise 2: Import the `jax` package**","metadata":{"papermill":{"duration":0.01668,"end_time":"2022-01-09T11:47:57.833721","exception":false,"start_time":"2022-01-09T11:47:57.817041","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import jax","metadata":{"papermill":{"duration":0.024826,"end_time":"2022-01-09T11:47:57.875624","exception":false,"start_time":"2022-01-09T11:47:57.850798","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.058248Z","iopub.execute_input":"2022-02-04T05:42:23.058882Z","iopub.status.idle":"2022-02-04T05:42:23.064382Z","shell.execute_reply.started":"2022-02-04T05:42:23.058821Z","shell.execute_reply":"2022-02-04T05:42:23.063272Z"},"trusted":true},"execution_count":4,"outputs":[]},{"cell_type":"markdown","source":"**Exercise 3: Display the version of `jax`**","metadata":{"papermill":{"duration":0.016778,"end_time":"2022-01-09T11:47:57.909635","exception":false,"start_time":"2022-01-09T11:47:57.892857","status":"completed"},"tags":[]}},{"cell_type":"code","source":"jax.__version__","metadata":{"papermill":{"duration":0.028537,"end_time":"2022-01-09T11:47:57.955199","exception":false,"start_time":"2022-01-09T11:47:57.926662","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.067747Z","iopub.execute_input":"2022-02-04T05:42:23.068207Z","iopub.status.idle":"2022-02-04T05:42:23.080320Z","shell.execute_reply.started":"2022-02-04T05:42:23.068163Z","shell.execute_reply":"2022-02-04T05:42:23.079156Z"},"trusted":true},"execution_count":5,"outputs":[{"execution_count":5,"output_type":"execute_result","data":{"text/plain":"'0.2.19'"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 4: Display the default backend of `jax`**","metadata":{"papermill":{"duration":0.016987,"end_time":"2022-01-09T11:47:57.989727","exception":false,"start_time":"2022-01-09T11:47:57.97274","status":"completed"},"tags":[]}},{"cell_type":"code","source":"jax.default_backend()","metadata":{"papermill":{"duration":0.050031,"end_time":"2022-01-09T11:47:58.057205","exception":false,"start_time":"2022-01-09T11:47:58.007174","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.082348Z","iopub.execute_input":"2022-02-04T05:42:23.082809Z","iopub.status.idle":"2022-02-04T05:42:23.095233Z","shell.execute_reply.started":"2022-02-04T05:42:23.082766Z","shell.execute_reply":"2022-02-04T05:42:23.093998Z"},"trusted":true},"execution_count":6,"outputs":[{"execution_count":6,"output_type":"execute_result","data":{"text/plain":"'tpu'"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 5: Display the devices of the backend**","metadata":{"papermill":{"duration":0.017468,"end_time":"2022-01-09T11:47:58.092399","exception":false,"start_time":"2022-01-09T11:47:58.074931","status":"completed"},"tags":[]}},{"cell_type":"code","source":"jax.devices()","metadata":{"papermill":{"duration":0.025322,"end_time":"2022-01-09T11:47:58.135399","exception":false,"start_time":"2022-01-09T11:47:58.110077","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.097137Z","iopub.execute_input":"2022-02-04T05:42:23.097789Z","iopub.status.idle":"2022-02-04T05:42:23.108974Z","shell.execute_reply.started":"2022-02-04T05:42:23.097681Z","shell.execute_reply":"2022-02-04T05:42:23.107921Z"},"trusted":true},"execution_count":7,"outputs":[{"execution_count":7,"output_type":"execute_result","data":{"text/plain":"[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 6: Create a JAX DeviceArray with values [10, 1, 24] and assign it to `data`**","metadata":{"papermill":{"duration":0.017909,"end_time":"2022-01-09T11:47:58.171633","exception":false,"start_time":"2022-01-09T11:47:58.153724","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data = jax.numpy.array([10, 1, 24])\ndata","metadata":{"papermill":{"duration":0.241411,"end_time":"2022-01-09T11:47:58.431082","exception":false,"start_time":"2022-01-09T11:47:58.189671","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.110824Z","iopub.execute_input":"2022-02-04T05:42:23.111995Z","iopub.status.idle":"2022-02-04T05:42:23.136424Z","shell.execute_reply.started":"2022-02-04T05:42:23.111938Z","shell.execute_reply":"2022-02-04T05:42:23.135343Z"},"trusted":true},"execution_count":8,"outputs":[{"execution_count":8,"output_type":"execute_result","data":{"text/plain":"DeviceArray([10,  1, 24], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 7: Display the type of `data`**","metadata":{"papermill":{"duration":0.018122,"end_time":"2022-01-09T11:47:58.467996","exception":false,"start_time":"2022-01-09T11:47:58.449874","status":"completed"},"tags":[]}},{"cell_type":"code","source":"type(data)","metadata":{"papermill":{"duration":0.027493,"end_time":"2022-01-09T11:47:58.513999","exception":false,"start_time":"2022-01-09T11:47:58.486506","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.138099Z","iopub.execute_input":"2022-02-04T05:42:23.138845Z","iopub.status.idle":"2022-02-04T05:42:23.146703Z","shell.execute_reply.started":"2022-02-04T05:42:23.138797Z","shell.execute_reply":"2022-02-04T05:42:23.145435Z"},"trusted":true},"execution_count":9,"outputs":[{"execution_count":9,"output_type":"execute_result","data":{"text/plain":"jax.interpreters.xla._DeviceArray"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 8: Display the shape of `data`**","metadata":{"papermill":{"duration":0.018392,"end_time":"2022-01-09T11:47:58.55108","exception":false,"start_time":"2022-01-09T11:47:58.532688","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data.shape","metadata":{"papermill":{"duration":0.02724,"end_time":"2022-01-09T11:47:58.598463","exception":false,"start_time":"2022-01-09T11:47:58.571223","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.148247Z","iopub.execute_input":"2022-02-04T05:42:23.148771Z","iopub.status.idle":"2022-02-04T05:42:23.160951Z","shell.execute_reply.started":"2022-02-04T05:42:23.148727Z","shell.execute_reply":"2022-02-04T05:42:23.159758Z"},"trusted":true},"execution_count":10,"outputs":[{"execution_count":10,"output_type":"execute_result","data":{"text/plain":"(3,)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 9: Transfer `data` to host and assign it to `data_host`**","metadata":{"papermill":{"duration":0.018859,"end_time":"2022-01-09T11:47:58.636233","exception":false,"start_time":"2022-01-09T11:47:58.617374","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data_host = jax.device_get(data)\ndata_host","metadata":{"papermill":{"duration":0.028099,"end_time":"2022-01-09T11:47:58.683844","exception":false,"start_time":"2022-01-09T11:47:58.655745","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.162882Z","iopub.execute_input":"2022-02-04T05:42:23.163442Z","iopub.status.idle":"2022-02-04T05:42:23.175198Z","shell.execute_reply.started":"2022-02-04T05:42:23.163390Z","shell.execute_reply":"2022-02-04T05:42:23.173834Z"},"trusted":true},"execution_count":11,"outputs":[{"execution_count":11,"output_type":"execute_result","data":{"text/plain":"array([10,  1, 24], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 10: Transfer `data_host` to device and assign it to `data_device`**","metadata":{"papermill":{"duration":0.019269,"end_time":"2022-01-09T11:47:58.722819","exception":false,"start_time":"2022-01-09T11:47:58.70355","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data_device = jax.device_put(data_host)\ndata_device","metadata":{"papermill":{"duration":0.028953,"end_time":"2022-01-09T11:47:58.771073","exception":false,"start_time":"2022-01-09T11:47:58.74212","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-04T05:42:23.177217Z","iopub.execute_input":"2022-02-04T05:42:23.177803Z","iopub.status.idle":"2022-02-04T05:42:23.193317Z","shell.execute_reply.started":"2022-02-04T05:42:23.177754Z","shell.execute_reply":"2022-02-04T05:42:23.192193Z"},"trusted":true},"execution_count":12,"outputs":[{"execution_count":12,"output_type":"execute_result","data":{"text/plain":"DeviceArray([10,  1, 24], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"<center>\n    This completes Set 1: JAX Introduction (Solutions 1-10) 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.019317,"end_time":"2022-01-09T11:47:58.810595","exception":false,"start_time":"2022-01-09T11:47:58.791278","status":"completed"},"tags":[]}}]}