{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Random Numbers And Arrays"
   ],
   "id": "site-title"
  },
  {
   "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": 2,
   "id": "9929fe42",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6394267984578837"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "random.seed(42)\n",
    "random.random()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e56e1f99",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random.randrange(1, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6fea7027",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 4,
     "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. The primary goal behind Jax's implementaition was more to have a deterministic state for the PRNG among multiple devices (say GPUs), so that the results can be easily reproduced. \n",
    "\n",
    "Usually, in a distributed setup (with multiple devices) 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). This means that one, you need to parallely generate random numbers[2] while keeping the state of your PRNG consistent and two, you need a way to share the same state between your devices. Jax shares PRNG state via key splitting[1]. You can read more from the sources listed below.\n",
    "\n",
    "**Sources:**\n",
    "1. [Splittable pseudorandom number generators using cryptographic hashing](https://dl.acm.org/doi/10.1145/2578854.2503784)\n",
    "2. [Parallel Random Numbers: As Easy as 1, 2, 3](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf)"
   ]
  },
  {
   "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": 5,
   "id": "2d2a2dc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array((), dtype=key<fry>) overlaying:\n",
       "[ 0 42]"
      ]
     },
     "execution_count": 5,
     "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": 6,
   "id": "628f6f48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.02830462, dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aeefa869",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.02830462, dtype=float32)"
      ]
     },
     "execution_count": 7,
     "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",
    "Because it's the same seed? If you want your random numbers to be reproducible, you should always have the same seed. \n",
    "\n",
    "To get a new random number, you need a new key and for that, you just have to __split__ the key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c90303a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "newkey, subkey = J.random.split(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4ae4fcb8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.07592554, dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J.random.normal(newkey)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "35ee4c54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.07592554, dtype=float32)"
      ]
     },
     "execution_count": 12,
     "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": 13,
   "id": "74191fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "newkey, *subkeys = J.random.split(key, 9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6d0c2b33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Array((), dtype=key<fry>) overlaying:\n",
      "[  64467757 2916123636]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2465931498  255383827]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[3134548294  894150801]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2954079971 3276725750]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2765691542  824333390]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2768684296 3055579793]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[2547012911 1371500959]\n",
      "Array((), dtype=key<fry>) overlaying:\n",
      "[1016697191 2390192106]\n"
     ]
    }
   ],
   "source": [
    "for sk in subkeys:\n",
    "    print(sk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "89b8829c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.60576403\n",
      "0.4323065\n",
      "-0.2818947\n",
      "0.6549178\n",
      "-0.2166012\n",
      "-0.25440374\n",
      "0.2886397\n",
      "0.14384735\n"
     ]
    }
   ],
   "source": [
    "for sk in subkeys:\n",
    "    print(J.random.normal(sk))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e9f82e9",
   "metadata": {},
   "source": [
    "### Randomized Vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "79535d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = J.random.key(99) # why should 42 get all the attention?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "381e31c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1.3088475 ],\n",
       "       [-0.25337982],\n",
       "       [-0.75657994],\n",
       "       [ 0.7655394 ],\n",
       "       [ 0.34965774],\n",
       "       [ 1.6598269 ],\n",
       "       [ 1.5128751 ],\n",
       "       [-0.79486036],\n",
       "       [-0.47534505],\n",
       "       [-1.8532473 ]], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = J.random.normal(key, shape=(10, 1))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1892271a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1.3088475 , -0.25337982, -0.75657994,  0.7655394 ],\n",
       "       [ 0.34965774,  1.6598269 ,  1.5128751 , -0.79486036],\n",
       "       [-0.47534505, -1.8532473 , -0.3786292 , -1.6755378 ],\n",
       "       [-0.9703495 ,  0.6404915 ,  0.58516824, -0.1952716 ],\n",
       "       [ 0.9720966 ,  0.3177939 ,  1.2931043 , -0.26136768],\n",
       "       [-1.5360676 , -0.921658  , -1.0590659 , -1.4680037 ]],      dtype=float32)"
      ]
     },
     "execution_count": 18,
     "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": 19,
   "id": "e38cb79e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 19,
     "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": 20,
   "id": "3ab236e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 2,  4,  6,  8, 10, 12, 14, 16, 18], dtype=int32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = b * 2\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "89f118ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.5, 1. , 1.5, 2. , 2.5, 3. , 3.5, 4. , 4.5], dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6e50018d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b + 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c829053c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1,  2,  3,  4,  5,  6,  7,  8,  9],\n",
       "       [ 2,  4,  6,  8, 10, 12, 14, 16, 18]], dtype=int32)"
      ]
     },
     "execution_count": 23,
     "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": 24,
   "id": "58cc8cc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3, dtype=int32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "cc9d666a",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[25]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mb\u001b[49m\u001b[43m[\u001b[49m\u001b[32;43m2\u001b[39;49m\u001b[43m]\u001b[49m = \u001b[32m99\u001b[39m\n\u001b[32m      2\u001b[39m b\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:617\u001b[39m, in \u001b[36m_unimplemented_setitem\u001b[39m\u001b[34m(self, i, x)\u001b[39m\n\u001b[32m    613\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_unimplemented_setitem\u001b[39m(\u001b[38;5;28mself\u001b[39m, i, x):\n\u001b[32m    614\u001b[39m   msg = (\u001b[33m\"\u001b[39m\u001b[33mJAX arrays are immutable and do not support in-place item assignment.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    615\u001b[39m          \u001b[33m\"\u001b[39m\u001b[33m Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method:\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    616\u001b[39m          \u001b[33m\"\u001b[39m\u001b[33m https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m--> \u001b[39m\u001b[32m617\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg.format(\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)))\n",
      "\u001b[31mTypeError\u001b[39m: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/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": 26,
   "id": "60033b36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 1,  2, 99,  4,  5,  6,  7,  8,  9], dtype=int32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.at[2].set(99)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "224c6be9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# doesn't change!\n",
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8052c5",
   "metadata": {},
   "source": [
    "But if I do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f8b59adc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 1,  2, 99,  4,  5,  6,  7,  8,  9], dtype=int32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_b = b.at[2].set(99)\n",
    "new_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d48ebe19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b"
   ]
  },
  {
   "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": 30,
   "id": "1060892e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)"
      ]
     },
     "execution_count": 30,
     "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": ".venv",
   "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.12.11"
  },
  "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
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}