notebooks/02_data_operations_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.015838,"end_time":"2022-02-04T05:32:28.988637","exception":false,"start_time":"2022-02-04T05:32:28.972799","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"<center>\n    This is Set 2: Data Operations (Solutions 11-20) 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.012769,"end_time":"2022-02-04T05:32:29.015298","exception":false,"start_time":"2022-02-04T05:32:29.002529","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.","metadata":{"papermill":{"duration":0.012743,"end_time":"2022-02-04T05:32:29.041006","exception":false,"start_time":"2022-02-04T05:32:29.028263","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!python3 -m pip install jax","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":8.902788,"end_time":"2022-02-04T05:32:37.956623","exception":false,"start_time":"2022-02-04T05:32:29.053835","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:17:33.706807Z","iopub.execute_input":"2022-02-07T06:17:33.707140Z","iopub.status.idle":"2022-02-07T06:17:43.322177Z","shell.execute_reply.started":"2022-02-07T06:17:33.707058Z","shell.execute_reply":"2022-02-07T06:17:43.321380Z"},"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: absl-py in /opt/conda/lib/python3.7/site-packages (from jax) (0.12.0)\nRequirement already satisfied: numpy>=1.18 in /opt/conda/lib/python3.7/site-packages (from jax) (1.19.5)\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":"code","source":"## import packages\nimport jax\nimport jax.numpy as jnp\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.874154,"end_time":"2022-02-04T05:32:39.851424","exception":false,"start_time":"2022-02-04T05:32:37.97727","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:17:43.323633Z","iopub.execute_input":"2022-02-07T06:17:43.323881Z","iopub.status.idle":"2022-02-07T06:18:22.348735Z","shell.execute_reply.started":"2022-02-07T06:17:43.323851Z","shell.execute_reply":"2022-02-07T06:18:22.347694Z"},"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 11: Create a matrix with values [[10, 1, 24], [20, 15, 14]] and assign it to `data`**","metadata":{"papermill":{"duration":0.013624,"end_time":"2022-02-04T05:32:39.878818","exception":false,"start_time":"2022-02-04T05:32:39.865194","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data = jnp.array([[10, 1, 24], [20, 15, 14]])\ndata","metadata":{"papermill":{"duration":0.028225,"end_time":"2022-02-04T05:32:39.923002","exception":false,"start_time":"2022-02-04T05:32:39.894777","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:22.354643Z","iopub.execute_input":"2022-02-07T06:18:22.355264Z","iopub.status.idle":"2022-02-07T06:18:22.407466Z","shell.execute_reply.started":"2022-02-07T06:18:22.355225Z","shell.execute_reply":"2022-02-07T06:18:22.406335Z"},"trusted":true},"execution_count":3,"outputs":[{"execution_count":3,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[10,  1, 24],\n             [20, 15, 14]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 12: Assign the transpose of `data` to `dataT`**","metadata":{"papermill":{"duration":0.013743,"end_time":"2022-02-04T05:32:39.951264","exception":false,"start_time":"2022-02-04T05:32:39.937521","status":"completed"},"tags":[]}},{"cell_type":"code","source":"dataT = data.T\ndataT","metadata":{"papermill":{"duration":0.25684,"end_time":"2022-02-04T05:32:40.222213","exception":false,"start_time":"2022-02-04T05:32:39.965373","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:22.409012Z","iopub.execute_input":"2022-02-07T06:18:22.410018Z","iopub.status.idle":"2022-02-07T06:18:22.809315Z","shell.execute_reply.started":"2022-02-07T06:18:22.409979Z","shell.execute_reply":"2022-02-07T06:18:22.808314Z"},"trusted":true},"execution_count":4,"outputs":[{"execution_count":4,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[10, 20],\n             [ 1, 15],\n             [24, 14]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 13: Assign the element of `data` at index [0, 2] to `value`**","metadata":{"papermill":{"duration":0.014461,"end_time":"2022-02-04T05:32:40.25141","exception":false,"start_time":"2022-02-04T05:32:40.236949","status":"completed"},"tags":[]}},{"cell_type":"code","source":"value = data[0, 2]\nvalue","metadata":{"papermill":{"duration":0.099289,"end_time":"2022-02-04T05:32:40.365523","exception":false,"start_time":"2022-02-04T05:32:40.266234","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:22.811115Z","iopub.execute_input":"2022-02-07T06:18:22.811900Z","iopub.status.idle":"2022-02-07T06:18:23.064798Z","shell.execute_reply.started":"2022-02-07T06:18:22.811852Z","shell.execute_reply":"2022-02-07T06:18:23.063762Z"},"trusted":true},"execution_count":5,"outputs":[{"execution_count":5,"output_type":"execute_result","data":{"text/plain":"DeviceArray(24, dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 14: Update the value of `data` at index [1, 1] to `100`**","metadata":{"papermill":{"duration":0.014495,"end_time":"2022-02-04T05:32:40.395481","exception":false,"start_time":"2022-02-04T05:32:40.380986","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data = data.at[1, 1].set(100)\ndata","metadata":{"papermill":{"duration":0.070823,"end_time":"2022-02-04T05:32:40.480925","exception":false,"start_time":"2022-02-04T05:32:40.410102","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.066158Z","iopub.execute_input":"2022-02-07T06:18:23.066811Z","iopub.status.idle":"2022-02-07T06:18:23.144452Z","shell.execute_reply.started":"2022-02-07T06:18:23.066745Z","shell.execute_reply":"2022-02-07T06:18:23.143542Z"},"trusted":true},"execution_count":6,"outputs":[{"execution_count":6,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[ 10,   1,  24],\n             [ 20, 100,  14]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 15: Add `41` to the value of `data` at index [0, 0]**","metadata":{"papermill":{"duration":0.014636,"end_time":"2022-02-04T05:32:40.510737","exception":false,"start_time":"2022-02-04T05:32:40.496101","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data = data.at[0, 0].add(41)\ndata","metadata":{"papermill":{"duration":0.04927,"end_time":"2022-02-04T05:32:40.574766","exception":false,"start_time":"2022-02-04T05:32:40.525496","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.145768Z","iopub.execute_input":"2022-02-07T06:18:23.146711Z","iopub.status.idle":"2022-02-07T06:18:23.208962Z","shell.execute_reply.started":"2022-02-07T06:18:23.146662Z","shell.execute_reply":"2022-02-07T06:18:23.207918Z"},"trusted":true},"execution_count":7,"outputs":[{"execution_count":7,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[ 51,   1,  24],\n             [ 20, 100,  14]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 16: Calculate the minimum values over axis=1 and assign it to `mins`**","metadata":{"papermill":{"duration":0.014949,"end_time":"2022-02-04T05:32:40.605037","exception":false,"start_time":"2022-02-04T05:32:40.590088","status":"completed"},"tags":[]}},{"cell_type":"code","source":"mins = data.min(axis=1)\nmins","metadata":{"papermill":{"duration":0.045563,"end_time":"2022-02-04T05:32:40.665761","exception":false,"start_time":"2022-02-04T05:32:40.620198","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.210493Z","iopub.execute_input":"2022-02-07T06:18:23.210917Z","iopub.status.idle":"2022-02-07T06:18:23.238996Z","shell.execute_reply.started":"2022-02-07T06:18:23.210873Z","shell.execute_reply":"2022-02-07T06:18:23.237949Z"},"trusted":true},"execution_count":8,"outputs":[{"execution_count":8,"output_type":"execute_result","data":{"text/plain":"DeviceArray([ 1, 14], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 17: Select the first row of values of `data` and assign it to `data_select`**","metadata":{"papermill":{"duration":0.015089,"end_time":"2022-02-04T05:32:40.69668","exception":false,"start_time":"2022-02-04T05:32:40.681591","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data_select = data[0]\ndata_select","metadata":{"papermill":{"duration":0.038962,"end_time":"2022-02-04T05:32:40.751039","exception":false,"start_time":"2022-02-04T05:32:40.712077","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.242338Z","iopub.execute_input":"2022-02-07T06:18:23.242654Z","iopub.status.idle":"2022-02-07T06:18:23.293343Z","shell.execute_reply.started":"2022-02-07T06:18:23.242618Z","shell.execute_reply":"2022-02-07T06:18:23.292407Z"},"trusted":true},"execution_count":9,"outputs":[{"execution_count":9,"output_type":"execute_result","data":{"text/plain":"DeviceArray([51,  1, 24], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 18: Append the row `data_select` to `data`**","metadata":{"papermill":{"duration":0.015551,"end_time":"2022-02-04T05:32:40.783104","exception":false,"start_time":"2022-02-04T05:32:40.767553","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data = jnp.vstack([data, data_select])\ndata","metadata":{"papermill":{"duration":0.045961,"end_time":"2022-02-04T05:32:40.845032","exception":false,"start_time":"2022-02-04T05:32:40.799071","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.294704Z","iopub.execute_input":"2022-02-07T06:18:23.295622Z","iopub.status.idle":"2022-02-07T06:18:23.335338Z","shell.execute_reply.started":"2022-02-07T06:18:23.295573Z","shell.execute_reply":"2022-02-07T06:18:23.334219Z"},"trusted":true},"execution_count":10,"outputs":[{"execution_count":10,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[ 51,   1,  24],\n             [ 20, 100,  14],\n             [ 51,   1,  24]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 19: Multiply the matrices `data` and `dataT` and assign it to `data_prod`**","metadata":{"papermill":{"duration":0.016189,"end_time":"2022-02-04T05:32:40.877297","exception":false,"start_time":"2022-02-04T05:32:40.861108","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data_prod = jnp.dot(data, dataT)\ndata_prod","metadata":{"papermill":{"duration":0.069366,"end_time":"2022-02-04T05:32:40.963984","exception":false,"start_time":"2022-02-04T05:32:40.894618","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.336845Z","iopub.execute_input":"2022-02-07T06:18:23.337518Z","iopub.status.idle":"2022-02-07T06:18:23.367679Z","shell.execute_reply.started":"2022-02-07T06:18:23.337476Z","shell.execute_reply":"2022-02-07T06:18:23.366799Z"},"trusted":true},"execution_count":11,"outputs":[{"execution_count":11,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[1087, 1371],\n             [ 636, 2096],\n             [1087, 1371]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 20: Convert the dtype of `data_prod` to `float32`**","metadata":{"papermill":{"duration":0.022086,"end_time":"2022-02-04T05:32:41.006371","exception":false,"start_time":"2022-02-04T05:32:40.984285","status":"completed"},"tags":[]}},{"cell_type":"code","source":"data_prod = jnp.array(data_prod, dtype=jnp.float32)\ndata_prod","metadata":{"papermill":{"duration":0.047491,"end_time":"2022-02-04T05:32:41.071596","exception":false,"start_time":"2022-02-04T05:32:41.024105","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2022-02-07T06:18:23.369721Z","iopub.execute_input":"2022-02-07T06:18:23.370052Z","iopub.status.idle":"2022-02-07T06:18:23.393505Z","shell.execute_reply.started":"2022-02-07T06:18:23.370013Z","shell.execute_reply":"2022-02-07T06:18:23.392454Z"},"trusted":true},"execution_count":12,"outputs":[{"execution_count":12,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[1087., 1371.],\n             [ 636., 2096.],\n             [1087., 1371.]], dtype=float32)"},"metadata":{}}]},{"cell_type":"markdown","source":"<center>\n    This completes Set 2: Data Operations (Solutions 11-20) 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.016134,"end_time":"2022-02-04T05:32:41.104868","exception":false,"start_time":"2022-02-04T05:32:41.088734","status":"completed"},"tags":[]}}]}