notebooks/03_pseudorandom_numbers_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":{}},{"cell_type":"markdown","source":"<center>\n    This is Set 3: Pseudorandom Numbers (Solutions 21-30) 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":{}},{"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* A sample array `sample_data` will be used for the exercises.","metadata":{}},{"cell_type":"code","source":"!python3 -m pip install jax","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-02-08T12:27:35.395322Z","iopub.execute_input":"2022-02-08T12:27:35.395691Z","iopub.status.idle":"2022-02-08T12:27:44.994623Z","shell.execute_reply.started":"2022-02-08T12:27:35.395604Z","shell.execute_reply":"2022-02-08T12:27:44.993776Z"},"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: 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: 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 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,"execution":{"iopub.status.busy":"2022-02-08T12:27:44.997617Z","iopub.execute_input":"2022-02-08T12:27:44.997988Z","iopub.status.idle":"2022-02-08T12:28:32.633980Z","shell.execute_reply.started":"2022-02-08T12:27:44.997940Z","shell.execute_reply":"2022-02-08T12:28:32.632951Z"},"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":"code","source":"## sample data\nsample_data = jnp.array([10, 1, 24, 20, 15, 14])\nsample_data","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:32.635524Z","iopub.execute_input":"2022-02-08T12:28:32.636139Z","iopub.status.idle":"2022-02-08T12:28:32.652507Z","shell.execute_reply.started":"2022-02-08T12:28:32.636105Z","shell.execute_reply":"2022-02-08T12:28:32.651278Z"},"trusted":true},"execution_count":3,"outputs":[{"execution_count":3,"output_type":"execute_result","data":{"text/plain":"DeviceArray([10,  1, 24, 20, 15, 14], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 21: Create a pseudorandom number generator key with seed=100 and assign it to `key`**","metadata":{}},{"cell_type":"code","source":"key = jax.random.PRNGKey(100)\nkey","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:32.655412Z","iopub.execute_input":"2022-02-08T12:28:32.656253Z","iopub.status.idle":"2022-02-08T12:28:33.278113Z","shell.execute_reply.started":"2022-02-08T12:28:32.656178Z","shell.execute_reply":"2022-02-08T12:28:33.277128Z"},"trusted":true},"execution_count":4,"outputs":[{"execution_count":4,"output_type":"execute_result","data":{"text/plain":"DeviceArray([  0, 100], dtype=uint32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 22: Create a subkey from `key` and assign it to `subkey`**","metadata":{}},{"cell_type":"code","source":"key, subkey = jax.random.split(key)\nsubkey","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:33.279510Z","iopub.execute_input":"2022-02-08T12:28:33.279917Z","iopub.status.idle":"2022-02-08T12:28:33.576807Z","shell.execute_reply.started":"2022-02-08T12:28:33.279879Z","shell.execute_reply":"2022-02-08T12:28:33.575803Z"},"trusted":true},"execution_count":5,"outputs":[{"execution_count":5,"output_type":"execute_result","data":{"text/plain":"array([3011861781, 1867493174], dtype=uint32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 23: Split `key` into seven subkeys `key_1`, `key_2`, `key_3`, `key_4`, `key_5`, `key_6` and `key_7`**","metadata":{}},{"cell_type":"code","source":"key_1, key_2, key_3, key_4, key_5, key_6, key_7 = jax.random.split(key, num=7)\nkey_1","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:33.578380Z","iopub.execute_input":"2022-02-08T12:28:33.578708Z","iopub.status.idle":"2022-02-08T12:28:33.839625Z","shell.execute_reply.started":"2022-02-08T12:28:33.578666Z","shell.execute_reply":"2022-02-08T12:28:33.838663Z"},"trusted":true},"execution_count":6,"outputs":[{"execution_count":6,"output_type":"execute_result","data":{"text/plain":"array([ 402730500, 1595431526], dtype=uint32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 24: Create a random permutation of `sample_data` using `key_1` and assign it to `data_permutation`**","metadata":{}},{"cell_type":"code","source":"data_permutation = jax.random.permutation(key_1, sample_data)\ndata_permutation","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:33.841028Z","iopub.execute_input":"2022-02-08T12:28:33.841646Z","iopub.status.idle":"2022-02-08T12:28:34.345059Z","shell.execute_reply.started":"2022-02-08T12:28:33.841602Z","shell.execute_reply":"2022-02-08T12:28:34.343656Z"},"trusted":true},"execution_count":7,"outputs":[{"execution_count":7,"output_type":"execute_result","data":{"text/plain":"DeviceArray([20, 14,  1, 10, 24, 15], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 25: Choose a random element from `sample_data` using `key_2` and assign it to `random_selection`**","metadata":{}},{"cell_type":"code","source":"random_selection = jax.random.choice(key_2, sample_data)\nrandom_selection","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:34.346864Z","iopub.execute_input":"2022-02-08T12:28:34.347610Z","iopub.status.idle":"2022-02-08T12:28:35.254576Z","shell.execute_reply.started":"2022-02-08T12:28:34.347557Z","shell.execute_reply":"2022-02-08T12:28:35.253743Z"},"trusted":true},"execution_count":8,"outputs":[{"execution_count":8,"output_type":"execute_result","data":{"text/plain":"DeviceArray(1, dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 26: Sample an integer between 10 and 24 using `key_3` and assign it to `sample_int`**","metadata":{}},{"cell_type":"code","source":"sample_int = jax.random.randint(key_3, shape=(1,), minval=10, maxval=24)\nsample_int","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:35.256157Z","iopub.execute_input":"2022-02-08T12:28:35.256806Z","iopub.status.idle":"2022-02-08T12:28:36.087709Z","shell.execute_reply.started":"2022-02-08T12:28:35.256764Z","shell.execute_reply":"2022-02-08T12:28:36.086371Z"},"trusted":true},"execution_count":9,"outputs":[{"execution_count":9,"output_type":"execute_result","data":{"text/plain":"DeviceArray([14], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 27: Sample two values from uniform distribution between 1 and 2 using `key_4` and assign it to `sample_uniform`**","metadata":{}},{"cell_type":"code","source":"sample_uniform = jax.random.uniform(key_4, shape=(2,), minval=1, maxval=2)\nsample_uniform","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:36.090799Z","iopub.execute_input":"2022-02-08T12:28:36.091681Z","iopub.status.idle":"2022-02-08T12:28:36.353448Z","shell.execute_reply.started":"2022-02-08T12:28:36.091633Z","shell.execute_reply":"2022-02-08T12:28:36.352348Z"},"trusted":true},"execution_count":10,"outputs":[{"execution_count":10,"output_type":"execute_result","data":{"text/plain":"DeviceArray([1.6274643, 1.1133162], dtype=float32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 28: Sample three values from bernoulli distribution using `key_5` and assign it to `sample_bernoulli`**","metadata":{}},{"cell_type":"code","source":"sample_bernoulli = jax.random.bernoulli(key_5, shape=(3,))\nsample_bernoulli","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:36.358031Z","iopub.execute_input":"2022-02-08T12:28:36.358715Z","iopub.status.idle":"2022-02-08T12:28:36.663860Z","shell.execute_reply.started":"2022-02-08T12:28:36.358655Z","shell.execute_reply":"2022-02-08T12:28:36.662872Z"},"trusted":true},"execution_count":11,"outputs":[{"execution_count":11,"output_type":"execute_result","data":{"text/plain":"DeviceArray([False,  True,  True], dtype=bool)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 29: Sample a 2x3 matrix from poisson distribution with λ=100 using `key_6` and assign it to `sample_poisson`**","metadata":{}},{"cell_type":"code","source":"sample_poisson = jax.random.poisson(key_6, shape=(2, 3), lam=100)\nsample_poisson","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:36.665346Z","iopub.execute_input":"2022-02-08T12:28:36.665775Z","iopub.status.idle":"2022-02-08T12:28:38.463111Z","shell.execute_reply.started":"2022-02-08T12:28:36.665731Z","shell.execute_reply":"2022-02-08T12:28:38.461993Z"},"trusted":true},"execution_count":12,"outputs":[{"execution_count":12,"output_type":"execute_result","data":{"text/plain":"DeviceArray([[ 88,  82, 110],\n             [ 89,  85,  98]], dtype=int32)"},"metadata":{}}]},{"cell_type":"markdown","source":"**Exercise 30: Sample a 2x3x4 array from normal distribution using `key_7` and assign it to `sample_normal`**","metadata":{}},{"cell_type":"code","source":"sample_normal = jax.random.normal(key_7, shape=(2, 3, 4))\nsample_normal","metadata":{"execution":{"iopub.status.busy":"2022-02-08T12:28:38.464716Z","iopub.execute_input":"2022-02-08T12:28:38.465035Z","iopub.status.idle":"2022-02-08T12:28:38.772551Z","shell.execute_reply.started":"2022-02-08T12:28:38.464993Z","shell.execute_reply":"2022-02-08T12:28:38.771632Z"},"trusted":true},"execution_count":13,"outputs":[{"execution_count":13,"output_type":"execute_result","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)"},"metadata":{}}]},{"cell_type":"markdown","source":"<center>\n    This completes Set 3: Pseudorandom Numbers (Solutions 21-30) 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":{}}]}