04_random_numbers_and_arrays.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cba72399",
   "metadata": {},
   "source": [
    "### Random numbers\n",
    "\n",
    "How important are random numbers when it comes to machine learning or deep learning? Well helluva important. How else are you going to initialize your parameters then? \n",
    "\n",
    "Actually a theoratical discussion on parameter initialization will hugely off track from the goal of these notebooks, so how about you read these wonderful pieces? (defintely later, because if you start reading them now and get lost you won't be able to get your interest back here).\n",
    "\n",
    "1. [Initializing neural networks](https://www.deeplearning.ai/ai-notes/initialization/)\n",
    "2. [Weight Initialization in Neural Networks: A Journey From the Basics to Kaiming](https://towardsdatascience.com/weight-initialization-in-neural-networks-a-journey-from-the-basics-to-kaiming-954fb9b47c79)\n",
    "3. [Weight Initialization for Deep Learning Neural Networks](https://machinelearningmastery.com/weight-initialization-for-deep-learning-neural-networks/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03799504",
   "metadata": {},
   "source": [
    "### Typical Random Numbers in python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9929fe42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6394267984578837"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "random.seed(42)\n",
    "random.random()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e56e1f99",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random.randrange(1, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6fea7027",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random.randrange(1, 10, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fbd6f4c",
   "metadata": {},
   "source": [
    "### Random in Jax\n",
    "\n",
    "Jax uses a different Pseudo Random Number Generator or, PRNG compared to the default in python or numpy. You can read more [here](https://github.com/google/jax/blob/main/design_notes/prng.md). The primary goal behind Jax's implementaition was more to the state of the PRNG consistent among multiple devices (say GPUs). \n",
    "\n",
    "What happens in a multi device setting is that the state of a PRNG changes (even for the same seed), which can lead to inconsistent results between the devices. Ideally you would want all your devices to sync the same PRNG configuration and produce random numbers the same way (i.e. without duplicates).\n",
    "\n",
    "Obviously there were other design considerations, which they have mentioned in the link above.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4305ed5f",
   "metadata": {},
   "source": [
    "So in Jax, you generate a key for a random number, which is similar to what you call a seed for a random generator. (If anybody knows why 42 is a commonly used seed, kindly let me know!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2d2a2dc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 0, 42], dtype=uint32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax as J\n",
    "\n",
    "key = J.random.key(42) # holy 42!\n",
    "key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "628f6f48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(-0.18471177, dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aeefa869",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(-0.18471177, dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3be15114",
   "metadata": {},
   "source": [
    "#### Whoa there! why the same number again? \n",
    "\n",
    "Weird isn't it! The thing is, random generators in jax don't keep track whethere a key was previously used or not. So everytime you want to generate a new random number, you have to create a new key. This is kinda weird. \n",
    "\n",
    "But you can get around this \"creating a new key every time\". You just have to __split__ the key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c90303a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "newkey, subkey = J.random.split(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ae4fcb8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.13790321, dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(newkey)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "35ee4c54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.13790321, dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(newkey) # one more time ....."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b416f6d6",
   "metadata": {},
   "source": [
    "### Multiple Splits\n",
    "\n",
    "![what if](./images/what_if.jpg)\n",
    "\n",
    "\n",
    "Off course! Legit question. You can mention how many keys you need with `split()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "74191fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "newkey, *subkeys = J.random.split(key, 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6d0c2b33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2015457675 1759218286]\n",
      "[ 372461012 2067635993]\n",
      "[1901925373 4124313125]\n",
      "[ 163989101 2463678267]\n",
      "[2310818676 2815820454]\n",
      "[3589763416 3424932669]\n",
      "[4221803610 1067686388]\n",
      "[2667717255 2997884537]\n"
     ]
    }
   ],
   "source": [
    "for sk in subkeys:\n",
    "    print(sk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "89b8829c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.43648595\n",
      "-2.2741115\n",
      "1.0372064\n",
      "-0.41117954\n",
      "-0.21331418\n",
      "0.9359986\n",
      "-0.3840317\n",
      "-0.92283344\n"
     ]
    }
   ],
   "source": [
    "for sk in subkeys:\n",
    "    print(J.random.normal(sk))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e9f82e9",
   "metadata": {},
   "source": [
    "### Randomized Vectors\n",
    "\n",
    "This was also, a reason for Jax's random generator design. Now why do you need vectors, you may ask? A bit of recalling, algorithms in machine learning require vectors and matrices ;)\n",
    "\n",
    "So far we have been writing hard coded matrices and vectors in these notebooks. This ends now. We use RANDOM!!!!!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "79535d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = J.random.key(99) # why should 42 get all the attention?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "381e31c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[-0.9513214 ],\n",
       "             [ 1.7788743 ],\n",
       "             [-1.4580659 ],\n",
       "             [-0.33244848],\n",
       "             [ 0.24782881],\n",
       "             [ 1.0122505 ],\n",
       "             [ 0.7675285 ],\n",
       "             [ 0.3141468 ],\n",
       "             [-0.01992193],\n",
       "             [ 0.42590362]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = J.random.normal(key, shape=(10, 1))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1892271a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 0.48635894, -1.0345986 , -1.1348429 ,  1.4427332 ],\n",
       "             [-0.34284604,  1.4199504 ,  1.0333503 , -0.92543155],\n",
       "             [-1.3507137 , -0.43571466, -0.21265179, -0.7674866 ],\n",
       "             [ 0.05622242,  0.16131395, -1.1077237 ,  0.71504277],\n",
       "             [-0.26486132,  0.16493076, -0.4747249 , -0.07925941],\n",
       "             [ 0.06515129, -0.42943704,  0.10599458, -1.6549199 ]],            dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mat_a = J.random.normal(key, shape=(6, 4))\n",
    "mat_a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0858d5c3",
   "metadata": {},
   "source": [
    "### Arrays in Jax\n",
    "\n",
    "If you're coming fresh out of a beginner python crash course, arrays are what you've known as lists, but a bit faster, better, more memory efficient etc. If you've experience working with numpy, Jax is just basic numpy with some magic sauce. Same applies for the arrays. Jax arrays are numpy arrays. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e38cb79e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# should we hard code again? nein!\n",
    "\n",
    "import jax.numpy as jnp\n",
    "\n",
    "b = jnp.arange(1, 10)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3ab236e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 2,  4,  6,  8, 10, 12, 14, 16, 18], dtype=int32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = b * 2\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "89f118ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([0.5, 1. , 1.5, 2. , 2.5, 3. , 3.5, 4. , 4.5], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6e50018d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b + 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c829053c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 1,  2,  3,  4,  5,  6,  7,  8,  9],\n",
       "             [ 2,  4,  6,  8, 10, 12, 14, 16, 18]], dtype=int32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_array = jnp.array(\n",
    "    [b, c]\n",
    ")\n",
    "new_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "dc789be3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 9)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_array.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e7f46d9",
   "metadata": {},
   "source": [
    "### Since they are like numpy arrays......\n",
    "\n",
    "You should be able to change a value at a specific index, right? Erm............"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "58cc8cc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(3, dtype=int32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "cc9d666a",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'<class 'jaxlib.xla_extension.DeviceArray'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [23], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m b[\u001b[39m2\u001b[39;49m] \u001b[39m=\u001b[39m \u001b[39m99\u001b[39m\n\u001b[1;32m      2\u001b[0m b\n",
      "File \u001b[0;32m~/Projects/jax_examples/venv/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:4946\u001b[0m, in \u001b[0;36m_unimplemented_setitem\u001b[0;34m(self, i, x)\u001b[0m\n\u001b[1;32m   4941\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_unimplemented_setitem\u001b[39m(\u001b[39mself\u001b[39m, i, x):\n\u001b[1;32m   4942\u001b[0m   msg \u001b[39m=\u001b[39m (\u001b[39m\"\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m object does not support item assignment. JAX arrays are \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   4943\u001b[0m          \u001b[39m\"\u001b[39m\u001b[39mimmutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   4944\u001b[0m          \u001b[39m\"\u001b[39m\u001b[39mor another .at[] method: \u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m   4945\u001b[0m          \u001b[39m\"\u001b[39m\u001b[39mhttps://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m-> 4946\u001b[0m   \u001b[39mraise\u001b[39;00m \u001b[39mTypeError\u001b[39;00m(msg\u001b[39m.\u001b[39mformat(\u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m)))\n",
      "\u001b[0;31mTypeError\u001b[0m: '<class 'jaxlib.xla_extension.DeviceArray'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html"
     ]
    }
   ],
   "source": [
    "b[2] = 99\n",
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "990922eb",
   "metadata": {},
   "source": [
    "JAX arrays are immutable. Which means they can't be modified once created. But you may need to modify values at one point in your wonderful experiments with glorious functions called AI. Jax allows that but differently. You can modify the array but it'll create a new array and preserve the older one. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "60033b36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 1,  2, 99,  4,  5,  6,  7,  8,  9], dtype=int32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.at[2].set(99)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "224c6be9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8052c5",
   "metadata": {},
   "source": [
    "But if I do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f8b59adc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([ 1,  2, 99,  4,  5,  6,  7,  8,  9], dtype=int32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_b = b.at[2].set(99)\n",
    "new_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d48ebe19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc43cd74",
   "metadata": {},
   "source": [
    "Now you have both the old and the new array. All this behaviour seems weird. How dare Jax won't let you do something you want? You've to wait a bit for that. Perhaps another notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62375044",
   "metadata": {},
   "source": [
    "### More weirdo\n",
    "\n",
    "```\n",
    "b.at[100].set(100)\n",
    "```\n",
    "\n",
    "Do you think this line will run? b only has a length of 9. Whatever you say I won't blame you. You're neither wrong or correct in the JAX world."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1060892e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.at[100].set(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7092cb1",
   "metadata": {},
   "source": [
    "JAX completely ignores out of bound errors and will keep your original array intact. I don't know why it works this way when JAX imposes so much strict restrictions on other areas. So keep track of your indexes in your projects otherwise you'll be scratching your head why some value didn't update. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "85f6f424c393e7c95d987d0a64d03300802113ad2045bd4cdb9a90d6d84115f4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}