Chapter 02 - Preliminaries (Tensors and Tensor Ops).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e54466b3",
   "metadata": {},
   "source": [
    "# Chapter 2: Preliminaries\n",
    "\n",
    "### Dhuvi Karthikeyan\n",
    "\n",
    "##### 1/09/2023"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "263264d0",
   "metadata": {},
   "source": [
    "**Updated: 01/11/23**: This notebook follows Chapter 2: Preliminaries from the Dive into Deep Learning Text (D2L) and serves as a space for notes for sections that I wanted to keep for future review. It doesn't contain all of the sections and while the names of the sections are slightly altered, I have tried to preserve the original hierarchical structure of the format and have deviated only where more intuitive for my mental model of the course material. This chapter serves as a getting your feet wet with deep learning frameworks a la creating tensors and familiarizing ourselves with their operations. I used this as a chance to test out JAX and compare it against to torch and tensorflow APIs for the same functions and test out their functionalities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e603430",
   "metadata": {},
   "source": [
    "## 2.1 Data Manipulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c3291a50",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-01-11 16:43:36.398663: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-01-11 16:43:37.013429: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-01-11 16:43:37.817330: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11.7/lib64:\n",
      "2023-01-11 16:43:37.817399: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda-11.7/lib64:\n",
      "2023-01-11 16:43:37.817405: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    }
   ],
   "source": [
    "# Import the Frameworks\n",
    "\n",
    "import torch\n",
    "import tensorflow as tf\n",
    "import jax\n",
    "from jax import numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "884566a0",
   "metadata": {},
   "source": [
    "### 2.1.1 Creating Tensors in the Three Frameworks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "73c15bb5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# PyTorch Instantiate a Tensor\n",
    "torchx = torch.arange(12, dtype=torch.float32)\n",
    "torchx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "03395b7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-01-11 16:43:39.487589: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:267] failed call to cuInit: CUDA_ERROR_UNKNOWN: unknown error\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(12,), dtype=float32, numpy=\n",
       "array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# TF Instantiate a Tensor\n",
    "tfx = tf.range(12, dtype=tf.float32)\n",
    "tfx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1a14bf96",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-01-11 16:43:39.571701: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:267] failed call to cuInit: CUDA_ERROR_UNKNOWN: unknown error\n",
      "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# JAX Instantiate a Tensor\n",
    "jaxx = jnp.arange(12)\n",
    "jaxx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd059084",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[[0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0.]],\n",
       "\n",
       "       [[0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0.]]], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.zeros((2, 3, 4)) #Also there's a JNP Ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "23cd56c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy=\n",
       "array([[[1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1.]],\n",
       "\n",
       "       [[1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1.],\n",
       "        [1., 1., 1., 1.]]], dtype=float32)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.ones((2, 3, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "73c40b5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0.]],\n",
       "\n",
       "        [[0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0.]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.zeros((2,3,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2804fc71",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3, 4), dtype=float32, numpy=\n",
       "array([[ 1.2134943 , -1.1871845 ,  0.18352774, -0.24870522],\n",
       "       [ 0.85458815,  0.2468478 ,  0.526479  ,  1.4140105 ],\n",
       "       [-0.14359604, -0.9495299 , -1.1702207 , -1.142411  ]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.random.normal(shape=[3, 4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "332d8f77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.8232,  1.1569, -0.2483,  1.8967],\n",
       "        [-0.7940,  2.3601, -0.5117, -1.4068],\n",
       "        [ 1.3473,  0.2974,  0.1927,  0.3635]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.randn(3,4) # Randn is not uniform! It is standard normal w/ mu = 0 and sd = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2c39e6dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1.1901639 , -1.0996888 ,  0.44367844,  0.5984697 ],\n",
       "       [-0.39189556,  0.69261974,  0.46018356, -2.068578  ],\n",
       "       [-0.21438177, -0.9898306 , -0.6789304 ,  0.27362573]],      dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Any call of a random function in JAX requires a key to be\n",
    "# specified, feeding the same key to a random function will\n",
    "# always result in the same sample being generated\n",
    "jax.random.normal(jax.random.PRNGKey(0), (3, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8aa84342",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[2, 1, 4, 3],\n",
       "        [1, 2, 3, 4],\n",
       "        [4, 3, 2, 1]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchY = torch.tensor([[2, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])\n",
    "torchY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fdc29bd6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[2, 1, 4, 3],\n",
       "       [1, 2, 3, 4],\n",
       "       [4, 3, 2, 1]], dtype=int32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxY = jnp.array([[2, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])\n",
    "jaxY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d13110d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3, 4), dtype=int32, numpy=\n",
       "array([[2, 1, 4, 3],\n",
       "       [1, 2, 3, 4],\n",
       "       [4, 3, 2, 1]], dtype=int32)>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfY = tf.constant([[2, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])\n",
    "tfY"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af5db713",
   "metadata": {},
   "source": [
    "### 2.1.2 Indexing and Slicing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6bac6e51",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchx.numel() == jaxx.size  # Number of Elements total (int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c4e85d75",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=int32, numpy=12>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.size(tfx) #Interesting output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ee973d80",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12]) (12,) (12,)\n"
     ]
    }
   ],
   "source": [
    "print(torchx.shape, jaxx.shape, tfx.shape) #All the same function call to get shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6d3835",
   "metadata": {},
   "source": [
    "#### Tensor Manipulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1ad6d61d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchX = torchx.reshape(3,4)\n",
    "torchX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "57a3aad4",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'torchx' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m### We can also infer a missing dimension by putting -1\u001b[39;00m\n\u001b[0;32m----> 3\u001b[0m \u001b[43mtorchx\u001b[49m\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m4\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'torchx' is not defined"
     ]
    }
   ],
   "source": [
    "### We can also infer a missing dimension by putting -1\n",
    "\n",
    "torchx.reshape(-1, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2f9ea5b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 0,  1,  2,  3],\n",
       "       [ 4,  5,  6,  7],\n",
       "       [ 8,  9, 10, 11]], dtype=int32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxX = jaxx.reshape(3,4)\n",
    "jaxX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "31f07c00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3, 4), dtype=float32, numpy=\n",
       "array([[ 0.,  1.,  2.,  3.],\n",
       "       [ 4.,  5.,  6.,  7.],\n",
       "       [ 8.,  9., 10., 11.]], dtype=float32)>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfX = tf.reshape(tfx, (3, 4))\n",
    "tfX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "445c0ccc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 8.,  9., 10., 11.]),\n",
       " tensor([[ 4.,  5.,  6.,  7.],\n",
       "         [ 8.,  9., 10., 11.]]))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchX[-1], torchX[1:3] # Indexing works much like numpy for all three frameworks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7864227a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([ 8,  9, 10, 11], dtype=int32),\n",
       " Array([[4, 5],\n",
       "        [8, 9]], dtype=int32))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxX[-1], jaxX[1:3, 0:2] # Notable when only one slice is specified it takes it from dim=0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31545366",
   "metadata": {},
   "source": [
    "**Tensor Mutability varies across the frameworks**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "69d2405d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[12., 12., 12., 12.],\n",
       "        [12., 12., 12., 12.],\n",
       "        [ 8.,  9., 10., 11.]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchX[:2, :] = 12\n",
    "torchX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "77cae094",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment",
     "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 tfX[\u001b[38;5;241m1\u001b[39m:\u001b[38;5;241m3\u001b[39m, :] \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m12\u001b[39m \u001b[38;5;66;03m#Tensors are immutable in TF and JAX\u001b[39;00m\n\u001b[1;32m      2\u001b[0m tfX\n",
      "\u001b[0;31mTypeError\u001b[0m: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment"
     ]
    }
   ],
   "source": [
    "tfX[1:3, :] = 12 #Tensors are immutable in TF and JAX\n",
    "tfX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "84443c1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'Variable:0' shape=(3, 4) dtype=float32, numpy=\n",
       "array([[ 0.,  1.,  2.,  3.],\n",
       "       [ 4.,  5.,  9.,  7.],\n",
       "       [ 8.,  9., 10., 11.]], dtype=float32)>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_var = tf.Variable(tfX)\n",
    "X_var[1, 2].assign(9)\n",
    "X_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4a6e3505",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 0,  1,  2,  3],\n",
       "       [ 4,  5, 17,  7],\n",
       "       [ 8,  9, 10, 11]], dtype=int32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# JAX arrays are immutable. `jax.numpy.ndarray.at` index\n",
    "# update operators create a new array with the corresponding\n",
    "# modifications made\n",
    "X_new_1 = jaxX.at[1, 2].set(17)\n",
    "X_new_1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "926dbe35",
   "metadata": {},
   "source": [
    "### 2.1.3 Tensor Ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "aad2bc84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([162754.7969, 162754.7969, 162754.7969, 162754.7969, 162754.7969,\n",
       "        162754.7969, 162754.7969, 162754.7969,   2980.9580,   8103.0840,\n",
       "         22026.4648,  59874.1406])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.exp(torchx) #Unary vs (binary) operators take map R -> R also analogous in JNP.exp and tf.exp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ef06a946",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[12., 12., 12., 12.],\n",
       "         [12., 12., 12., 12.],\n",
       "         [ 8.,  9., 10., 11.],\n",
       "         [ 2.,  1.,  4.,  3.],\n",
       "         [ 1.,  2.,  3.,  4.],\n",
       "         [ 4.,  3.,  2.,  1.]]),\n",
       " tensor([[12., 12., 12., 12.,  2.,  1.,  4.,  3.],\n",
       "         [12., 12., 12., 12.,  1.,  2.,  3.,  4.],\n",
       "         [ 8.,  9., 10., 11.,  4.,  3.,  2.,  1.]]))"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cat((torchX, torchY), dim=0), torch.cat((torchX, torchY), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "78c9a09b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11],\n",
       "        [ 2,  1,  4,  3],\n",
       "        [ 1,  2,  3,  4],\n",
       "        [ 4,  3,  2,  1]], dtype=int32),\n",
       " Array([[ 0,  1,  2,  3,  2,  1,  4,  3],\n",
       "        [ 4,  5,  6,  7,  1,  2,  3,  4],\n",
       "        [ 8,  9, 10, 11,  4,  3,  2,  1]], dtype=int32))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.concatenate((jaxX, jaxY), axis=0), jnp.concatenate((jaxX, jaxY), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7838d524",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor: shape=(6, 4), dtype=float32, numpy=\n",
       " array([[ 0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.],\n",
       "        [ 0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.]], dtype=float32)>,\n",
       " <tf.Tensor: shape=(3, 8), dtype=float32, numpy=\n",
       " array([[ 0.,  1.,  2.,  3.,  0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.,  4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.,  8.,  9., 10., 11.]], dtype=float32)>)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.concat([tfX, tfX], axis=0), tf.concat([tfX, tfX], axis=1) #tfY was dtype int32 and refused to concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "591d79cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[12., 12., 12., 12.],\n",
       "         [12., 12., 12., 12.],\n",
       "         [ 8.,  9., 10., 11.]]),\n",
       " tensor([32., 33., 34., 35.]),\n",
       " tensor([48., 48., 38.]),\n",
       " tensor(134.))"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchX, torchX.sum(dim=0), torchX.sum(dim=1), torchX.sum() #Reduce Sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0ce55f90",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]], dtype=int32),\n",
       " Array([12, 15, 18, 21], dtype=int32),\n",
       " Array([ 6, 22, 38], dtype=int32),\n",
       " Array(66, dtype=int32))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxX, jaxX.sum(axis=0), jaxX.sum(axis=1), jaxX.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "03e77c99",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor: shape=(3, 4), dtype=float32, numpy=\n",
       " array([[ 0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.]], dtype=float32)>,\n",
       " <tf.Tensor: shape=(4,), dtype=float32, numpy=array([12., 15., 18., 21.], dtype=float32)>,\n",
       " <tf.Tensor: shape=(3,), dtype=float32, numpy=array([ 6., 22., 38.], dtype=float32)>,\n",
       " <tf.Tensor: shape=(), dtype=float32, numpy=66.0>)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfX, tf.reduce_sum(tfX, axis=0), tf.reduce_sum(tfX, axis=1), tf.reduce_sum(tfX)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f4e657",
   "metadata": {},
   "source": [
    "### 2.1.4 Broadcasting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01a550e5",
   "metadata": {},
   "source": [
    "Broadcasting works the same way as in NumPy where if two arrays are of different dimension, try expanding one or both arrays by duplicating elements along axes with length 1 so they have the same shape at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2268e379",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[0],\n",
       "        [1],\n",
       "        [2]], dtype=int32),\n",
       " Array([[0, 1]], dtype=int32))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = jnp.arange(3).reshape((3, 1))\n",
    "b = jnp.arange(2).reshape((1, 2))\n",
    "a, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "16fe96ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0, 1],\n",
       "       [1, 2],\n",
       "       [2, 3]], dtype=int32)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beec2fbe",
   "metadata": {},
   "source": [
    "### 2.1.5 Saving Memory (In-Place Operations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "fabd891e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "140173843992464\n",
      "After the addition 140173445156592\n"
     ]
    }
   ],
   "source": [
    "print(id(torchY))   # Can be a nuisance if updates don't happen in place and old params are used\n",
    "torchY = torchY + torchX\n",
    "print(\"After the addition\", id(torchY))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "53468158",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id(Z): 140173445160832\n",
      "id(Z): 140173445160832\n"
     ]
    }
   ],
   "source": [
    "torchZ = torch.zeros_like(torchY)\n",
    "print('id(Z):', id(torchZ))\n",
    "torchZ[:] = torchX + torchY #TF has the tfT.assign(operation)\n",
    "print('id(Z):', id(torchZ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "fc542a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# JAX arrays do not allow in-place operations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "028a2950",
   "metadata": {},
   "source": [
    "### 2.1.6 Conversions to Other Python Objects\n",
    "\n",
    "This is handled inplace conveniently for us by the backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "610beb21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[False, False, False, False],\n",
       "       [False, False, False, False],\n",
       "       [ True,  True,  True,  True]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchX.numpy() == tfX.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "662a8665",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "numpy.ndarray"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(jax.device_get(jaxX)) #The opposite is device_put"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e573135",
   "metadata": {},
   "source": [
    "## 2.3 Linear Algebra"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84878ac4",
   "metadata": {},
   "source": [
    "### 2.3.3 Matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "30585b29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[14., 13., 12.],\n",
       "        [13., 14., 12.],\n",
       "        [16., 15., 12.],\n",
       "        [15., 16., 12.]])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchY.T #Transpose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "999cd610",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[2, 1, 4],\n",
       "       [1, 2, 3],\n",
       "       [4, 3, 2],\n",
       "       [3, 4, 1]], dtype=int32)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jaxY.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ebfcd241",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(4, 3), dtype=int32, numpy=\n",
       "array([[2, 1, 4],\n",
       "       [1, 2, 3],\n",
       "       [4, 3, 2],\n",
       "       [3, 4, 1]], dtype=int32)>"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.transpose(tfY)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45c96812",
   "metadata": {},
   "source": [
    "### 2.3.4 Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "600bf539",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.,  1.,  2.,  3.],\n",
       "         [ 4.,  5.,  6.,  7.],\n",
       "         [ 8.,  9., 10., 11.]],\n",
       "\n",
       "        [[12., 13., 14., 15.],\n",
       "         [16., 17., 18., 19.],\n",
       "         [20., 21., 22., 23.]]])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)\n",
    "torch3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "569b6825",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14, 15],\n",
       "        [16, 17, 18, 19],\n",
       "        [20, 21, 22, 23]]], dtype=int32)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax3d = jnp.arange(24).reshape(2, 3, 4)\n",
    "jax3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "c2293c9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy=\n",
       "array([[[ 0.,  1.,  2.,  3.],\n",
       "        [ 4.,  5.,  6.,  7.],\n",
       "        [ 8.,  9., 10., 11.]],\n",
       "\n",
       "       [[12., 13., 14., 15.],\n",
       "        [16., 17., 18., 19.],\n",
       "        [20., 21., 22., 23.]]], dtype=float32)>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf3d = tf.reshape(tf.range(24, dtype=tf.float32), (2,3,4))\n",
    "tf3d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ed513b6",
   "metadata": {},
   "source": [
    "### 2.3.5 Basic Properties of Tensor Arithmetic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9add857e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[  0.,   1.,   4.,   9.],\n",
       "         [ 16.,  25.,  36.,  49.],\n",
       "         [ 64.,  81., 100., 121.]],\n",
       "\n",
       "        [[144., 169., 196., 225.],\n",
       "         [256., 289., 324., 361.],\n",
       "         [400., 441., 484., 529.]]])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d * torch3d # Hadamard Product or Element-wise product (*) is used across frameworks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c8ab5ee9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[[ 0,  2,  4,  6],\n",
       "        [ 8, 10, 12, 14],\n",
       "        [16, 18, 20, 22]],\n",
       "\n",
       "       [[24, 26, 28, 30],\n",
       "        [32, 34, 36, 38],\n",
       "        [40, 42, 44, 46]]], dtype=int32)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax3d + jax3d  # Elementwise operator (+) is the same across frameworks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2031a8e1",
   "metadata": {},
   "source": [
    "### 2.3.6 Reduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "40fc3af0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(11.5000),\n",
       " Array([[ 6.,  7.,  8.,  9.],\n",
       "        [10., 11., 12., 13.],\n",
       "        [14., 15., 16., 17.]], dtype=float32),\n",
       " <tf.Tensor: shape=(2, 4), dtype=float32, numpy=\n",
       " array([[ 4.,  5.,  6.,  7.],\n",
       "        [16., 17., 18., 19.]], dtype=float32)>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d.mean(), jax3d.mean(axis=0), tf.reduce_mean(tf3d, axis=1) #Finding the mean"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "522d7425",
   "metadata": {},
   "source": [
    "### 2.3.7 Non-Reduction Sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "966c9e32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[12., 14., 16., 18.],\n",
       "          [20., 22., 24., 26.],\n",
       "          [28., 30., 32., 34.]]]),\n",
       " torch.Size([2, 1, 4]))"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d.sum(dim=0, keepdims=True), torch3d.sum(dim=1, keepdims=True).shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ccab5ba8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.,  1.,  2.,  3.],\n",
       "         [ 4.,  5.,  6.,  7.],\n",
       "         [ 8.,  9., 10., 11.]],\n",
       "\n",
       "        [[12., 13., 14., 15.],\n",
       "         [16., 17., 18., 19.],\n",
       "         [20., 21., 22., 23.]]])"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "3a66e885",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[12., 14., 16., 18.],\n",
       "         [20., 22., 24., 26.],\n",
       "         [28., 30., 32., 34.]]),\n",
       " tensor([[12., 15., 18., 21.],\n",
       "         [48., 51., 54., 57.]]),\n",
       " tensor([[ 6., 22., 38.],\n",
       "         [54., 70., 86.]]))"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d.sum(dim=0), torch3d.sum(dim=1), torch3d.sum(dim=2) #Adding across dim0, dim1, dim2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "38c5acd7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 0.,  1.,  2.,  3.],\n",
       "          [ 4.,  5.,  6.,  7.],\n",
       "          [ 8.,  9., 10., 11.]],\n",
       " \n",
       "         [[12., 14., 16., 18.],\n",
       "          [20., 22., 24., 26.],\n",
       "          [28., 30., 32., 34.]]]),\n",
       " tensor([[[ 0.,  1.,  2.,  3.],\n",
       "          [ 4.,  6.,  8., 10.],\n",
       "          [12., 15., 18., 21.]],\n",
       " \n",
       "         [[12., 13., 14., 15.],\n",
       "          [28., 30., 32., 34.],\n",
       "          [48., 51., 54., 57.]]]),\n",
       " tensor([[[ 0.,  1.,  3.,  6.],\n",
       "          [ 4.,  9., 15., 22.],\n",
       "          [ 8., 17., 27., 38.]],\n",
       " \n",
       "         [[12., 25., 39., 54.],\n",
       "          [16., 33., 51., 70.],\n",
       "          [20., 41., 63., 86.]]]))"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch3d.cumsum(dim=0), torch3d.cumsum(dim=1), torch3d.cumsum(dim=2) #CUMSUM across dim0, dim1, dim2 same in JAX, tf.cumsum "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b8b4b58",
   "metadata": {},
   "source": [
    "### 2.3.8 Dot Products"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71ad1895",
   "metadata": {},
   "source": [
    "torch.dot(torch.ones(3, dtype=torch.float32), torch.zeros(3, dtype=torch.float32))\n",
    "\n",
    "jnp.dot(jnp.ones(3, dtype=jnp.float32), jnp.zeros(3, dtype=jnp.float32))\n",
    "\n",
    "#tf.tensordot(tf.ones(3, dtype=tf.float32), tf.zeros(3, dtype=tf.float32), axes=1) #CUBLAS error"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "273abc4f",
   "metadata": {},
   "source": [
    "### 2.3.9 Matrix-Vector Products"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "6a8a9880",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([48., 48., 38.])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Torch is A@x\n",
    "# Jax is jnp.matmul(Z, x)\n",
    "# TF is separated tf.matvec and tf.matmul\n",
    "\n",
    "torchX@torch.ones(4, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c79655",
   "metadata": {},
   "source": [
    "### 2.3.10 Matrix-Matrix Products"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5a6d696",
   "metadata": {},
   "outputs": [],
   "source": [
    "torchX@torch.ones(3,4, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cac5af",
   "metadata": {},
   "source": [
    "### 2.3.11 Norms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0beb3ca",
   "metadata": {},
   "source": [
    "Norms all have the property of mapping vectors to scalars and have the following properties:\n",
    "\n",
    "1. $ ||ax|| = |a|||x|| $\n",
    "2. $ || x + y || \\leq ||x|| + ||y||$\n",
    "3. $ ||x|| >= 0 $ \n",
    "\n",
    "Lp Norm:\n",
    "\n",
    "$$ ||x||_p = (\\sum_{i=1}^n|x_i|^p)^{1/p} $$\n",
    "\n",
    "Frobenius Norm:\n",
    "\n",
    "$$||X||_F = \\sqrt{\\sum_i \\sum_j x_{ij}^2} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "a533f8c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Norms\n",
    "\n",
    "#L2 Norm = Euclidean Norm = Vector version of the Frobenius Norm\n",
    "#L1 Norm = Manhattan Distance = Sum of the Absolute Values of Vector Elements\n",
    "\n",
    "#Norms are subjectively easier to compute in JAX due to specification of order to wrapping func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "d903f7aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Torch and TF Norms\n",
    "\n",
    "#torch.norm(v) and tf.norm(u) are the Frobenius Norm also applies to the vectors (l2)\n",
    "#torch.abs(v).sum() and tf.reduce_sum(tf.abs(u)) is the L2 norm or the manhattan distance\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "ebc105ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(4., dtype=float32)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# JAX \n",
    "\n",
    "jnp.linalg.norm(jnp.ones((4,9)), ord=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "5e44c757",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(5.9999995, dtype=float32), Array(6., dtype=float32))"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.linalg.norm(jnp.ones((4,9)), ord=2), jnp.linalg.norm(jnp.ones((4,9)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6513b3bd",
   "metadata": {},
   "source": [
    "## 2.4 Calculus Refresher"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e13cd31",
   "metadata": {},
   "source": [
    "### 2.4.3 Partial Derivatives and Gradients\n",
    "\n",
    "* Gradient is the vector of partial derivatives\n",
    "* For all $A \\in \\mathbb{R^{m x n}} \\nabla Ax = A^T$\n",
    "* For all $A \\in \\mathbb{R^{m x n}} \\nabla x^TA = A$\n",
    "* For all square $A \\in \\mathbb{R^{n x n}} \\nabla x^TAx = (A+A^T)x$\n",
    "* $\\nabla ||x||^2 = \\nabla x^Tx = 2x$, also holds for Frobenius norm of a vector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29824050",
   "metadata": {},
   "source": [
    "### 2.4.4 Chain Rule\n",
    "\n",
    "$$ y = f(\\mathbf{u}) \\text{ where} u \\in \\mathbb{R}^{m}$$\n",
    "$$ \\mathbf{u} = g(\\mathbf{x}) \\text{ where} x \\in \\mathbb{R}^{n}$$ \n",
    "$$ \\frac{\\partial y}{\\partial x_i} = \\frac{\\partial y}{\\partial u_1}\\frac{\\partial u_1}{\\partial x_i} + \\frac{\\partial y}{\\partial u_2}\\frac{\\partial u_2}{\\partial x_i} ... + \\frac{\\partial y}{\\partial u_m}\\frac{\\partial u_m}{\\partial x_i}$$\n",
    "\n",
    "This is the same as stating succinctly that there exists an $\\mathbf{A} \\in \\mathbb{R}^{n x m}$:\n",
    "\n",
    "$$ \\nabla_x y = \\mathbf{A}\\nabla_u y$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb34c509",
   "metadata": {},
   "source": [
    "## 2.5 Automatic Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2d5b1e",
   "metadata": {},
   "source": [
    "### 2.5.1 A Simple Function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f3fa47",
   "metadata": {},
   "source": [
    "Auto calculate the gradients for the simple function: $y = 2x^2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "2a0d48d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(28., grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Torch\n",
    "\n",
    "torchz = torch.arange(4.0, dtype=torch.float32, requires_grad=True)\n",
    "torchz.grad #is None because backwards hasn't been called\n",
    "torchu = 2 * torch.dot(torchz, torchz)\n",
    "torchu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "0820a20a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dxk/miniconda3/lib/python3.10/site-packages/torch/autograd/__init__.py:197: UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:109.)\n",
      "  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([ 0.,  4.,  8., 12.])"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchu.backward()\n",
    "torchz.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "ff311d0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 1., 1., 1.])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# To calculate another function of x we need to first clear the gradients of torchz or clear the gradients\n",
    "\n",
    "torchz.grad.zero_()\n",
    "torchu = torchz.sum()\n",
    "torchu.backward()\n",
    "torchz.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "fda1b54b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(28., dtype=float32)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# JAX\n",
    "\n",
    "jaxz = jnp.arange(4.0) # Defaults to gradients turned on?\n",
    "jaxu = lambda x: 2 * jnp.dot(x,x) # Jax gradients are passed through the function\n",
    "jaxu(jaxz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "b15c23eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 0.,  4.,  8., 12.], dtype=float32)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax import grad\n",
    "\n",
    "jaxz_grad = grad(jaxu)(jaxz)\n",
    "jaxz_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "a020fcaf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 1., 1., 1.], dtype=float32)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Unlike Torch the gradients don't need to be cleared so running the following will give out the correct values\n",
    "\n",
    "jaxu = lambda x: x.sum()\n",
    "jaxz_grad = grad(jaxu)(jaxz)\n",
    "jaxz_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "07213881",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(4,), dtype=float32, numpy=array([ 0.,  4.,  8., 12.], dtype=float32)>"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Tensorflow\n",
    "\n",
    "tfz = tf.range(4, dtype=tf.float32)\n",
    "tfz = tf.Variable(tfz)                      #This code is necessary to store the grads as they are computed Container class\n",
    "with tf.GradientTape() as t:\n",
    "    tfu = 2 * tf.tensordot(tfz, tfz, axes=1) # Gives the tensor with scalar value 28.0\n",
    "t.gradient(tfu, tfz)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f201da",
   "metadata": {},
   "source": [
    "### 2.5.2 Backward for Non-Scalar Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b239a8",
   "metadata": {},
   "source": [
    "https://github.com/yang-zhang/yang-zhang.github.io/blob/master/ds_code/pytorch-backward-gradient-examples.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e8d2c5",
   "metadata": {},
   "source": [
    "The jacobian is a generalization of the gradient. The gradient is the vector of derivatives of a scalar valued output over a vector valued input. So for computing the loss, the gradient is what is being calculated. The jacobian on the other hand is what is used when an input of vectors maps to an output of vectors. It is the matrix of partial derivatives where the number of columns is equal to the dimension of the output and the number of rows is equal to the dimension of the input. Jacobians are seldom used in ML, instead using the sum of gradients w.r.t to a partical dimension of y."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "14092b70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0., 2., 4., 6.])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Torch\n",
    "torchz.grad.zero_()\n",
    "torchu = torchz * torchz\n",
    "torchu.backward(gradient=torch.ones(len(torchu)))  # Faster: y.sum().backward()\n",
    "torchz.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "2b1710e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 2., 4., 6.], dtype=float32)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#JAX\n",
    "ajax = lambda x: x * x\n",
    "# `grad` is only defined for scalar output functions\n",
    "grad(lambda x: ajax(x).sum())(jaxz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "0439c5ea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 2., 4., 6.], dtype=float32)>"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#TF\n",
    "with tf.GradientTape() as t:\n",
    "    tfu = tfz * tfz\n",
    "t.gradient(tfu, tfz)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e657fa",
   "metadata": {},
   "source": [
    "### 2.5.3 Detaching from Computational Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "d5a6291b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Torch': tensor([[14., 13., 16., 15.],\n",
       "         [13., 14., 15., 16.],\n",
       "         [12., 12., 12., 12.]]),\n",
       " 'TF': <tf.Tensor: shape=(3, 4), dtype=int32, numpy=\n",
       " array([[2, 1, 4, 3],\n",
       "        [1, 2, 3, 4],\n",
       "        [4, 3, 2, 1]], dtype=int32)>,\n",
       " 'JAX': Array([[2, 1, 4, 3],\n",
       "        [1, 2, 3, 4],\n",
       "        [4, 3, 2, 1]], dtype=int32)}"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{\n",
    "    'Torch': torchY.detach(),\n",
    "    'TF': tf.stop_gradient(tfY),\n",
    "    'JAX':jax.lax.stop_gradient(jaxY)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ad5ab15",
   "metadata": {},
   "source": [
    "### 2.5.4 Gradients and Python Control Flow\n",
    "\n",
    "Nontrivial point but gradients are able to be calculated various if conditions as long as there is a linear (smooth) flow of I/O through function transformations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf934bb9",
   "metadata": {},
   "source": [
    "## 2.6 Probability and Statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e6419a0",
   "metadata": {},
   "source": [
    "Conditional Independence: $P(A,B|C) = P(A|C) P(B|C)$\n",
    "\n",
    "$$ E[X+Y] = E[X] + E[Y] $$\n",
    "$$ Var[X] = E[(X-E[X])^2] = E[X^2] - E[X]^2 $$\n",
    "$$ Var[aX] = a^2Var[X] $$\n",
    "$$ Var[aX + bY] = a^2Var[X] + b^2Var[Y] + 2abCov(X,Y)$$\n",
    "\n",
    "Covariance matrix is calculated as: $E[(x-\\mu)(x-\\mu)^T]$\n",
    "\n",
    "* Aleatoric Uncertainty: Uncertainty intrinsic to the problem\n",
    "* Epistemic Uncertainty: Uncertainty over model params\n",
    "\n",
    "Chebyshev Inequality: \n",
    "\n",
    "$$ P(|X-\\mu| \\geq k\\sigma) \\leq \\frac{1}{k^2}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8eca6bb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}