{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.8.10 64-bit ('tf': conda)"
    },
    "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.8.10"
    },
    "colab": {
      "name": "Gradients_in_TF2.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "interpreter": {
      "hash": "859cadc33cf9c93d9354f622d0bfdc75f7d75fc23a9b22a18dfc646072acf83e"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "%matplotlib inline \n",
        "import matplotlib.pyplot as plt"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Init Plugin\n",
            "Init Graph Optimizer\n",
            "Init Kernel\n"
          ]
        }
      ],
      "metadata": {
        "id": "hMeZ2inBVTdF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Tensors (https://www.tensorflow.org/guide/tensor):\n",
        "### Tensors are multi-dimensional arrays with a uniform type, similar to numpy arrays\n",
        "\n"
      ],
      "metadata": {
        "id": "8O8igM3SydJD"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "source": [
        "tensor1 = tf.constant(4)\n",
        "print('\\ntensor1: ',tensor1)\n",
        "\n",
        "tensor2 = tf.constant([1.0,2.0,3.0])\n",
        "print('\\ntensor2: ',tensor2)\n",
        "\n",
        "tensor3 = tf.constant([[1.0,2.0,3.0],[3.4,1.0,1.0]])\n",
        "print('\\ntensor3: ',tensor3)\n",
        "\n",
        "tensor4 = tf.linspace(0,1,10)\n",
        "print('\\ntensor4: ',tensor4)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "tensor1: Metal device set to: Apple M1\n",
            " tf.Tensor(4, shape=(), dtype=int32)\n",
            "\n",
            "systemMemory: 16.00 GB\n",
            "maxCacheSize: 5.33 GB\n",
            "\n",
            "\n",
            "tensor2:  tf.Tensor([1. 2. 3.], shape=(3,), dtype=float32)\n",
            "\n",
            "tensor3:  tf.Tensor(\n",
            "[[1.  2.  3. ]\n",
            " [3.4 1.  1. ]], shape=(2, 3), dtype=float32)\n",
            "\n",
            "tensor4:  tf.Tensor(\n",
            "[0.         0.11111111 0.22222222 0.33333333 0.44444444 0.55555556\n",
            " 0.66666667 0.77777778 0.88888889 1.        ], shape=(10,), dtype=float64)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2021-09-29 08:48:08.666806: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
            "2021-09-29 08:48:08.667032: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
          ]
        }
      ],
      "metadata": {
        "id": "1vBO7fdAygTL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ed0f1e07-6912-4df8-b685-63a30682f9c5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Tensors can be converted to numpy arrays"
      ],
      "metadata": {
        "id": "v8V8W9Uf0KKR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "source": [
        "tensor4_np = tensor4.numpy()\n",
        "print('\\ntensor4 converted numpy array: ',tensor4_np)\n",
        "print('\\ntype: ',type(tensor4_np))"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "tensor4 converted numpy array:  [0.         0.11111111 0.22222222 0.33333333 0.44444444 0.55555556\n",
            " 0.66666667 0.77777778 0.88888889 1.        ]\n",
            "\n",
            "type:  <class 'numpy.ndarray'>\n"
          ]
        }
      ],
      "metadata": {
        "id": "fl5ut3D00Qy-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "53ed13b6-34f0-4b15-c1e4-50fcde631f4b"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "### All tensors are immutable like Python numbers and strings: you can never update the contents of a tensor, only create a new one.\n",
        "\n"
      ],
      "metadata": {
        "id": "6RVCxdcx0AaP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "source": [
        "print('6th element of tensor4: ',tensor4[5].numpy())\n",
        "tensor4[5] = 100.0"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "6th element of tensor4:  0.5555555555555556\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "TypeError",
          "evalue": "'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m/var/folders/6h/mmvp5df90fb3fsqtc5d9mxt40000gn/T/ipykernel_1503/2405235518.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'6th element of tensor4: '\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtensor4\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mtensor4\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m5\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m100.0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mTypeError\u001b[0m: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment"
          ]
        }
      ],
      "metadata": {
        "id": "fA8stmW10za6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 201
        },
        "outputId": "d4330068-03b2-4a67-8034-d2eae445623c"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Variables (https://www.tensorflow.org/guide/variable):\n",
        "\n",
        "### A variable is a tensor whose value can be changed by running operations on it"
      ],
      "metadata": {
        "id": "cZ5MAgDn1kMQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "var1 = tf.Variable(4)\n",
        "print('\\nvar1: ',var1)\n",
        "\n",
        "var2 = tf.Variable(tensor2)\n",
        "print('\\nvar2: ',var2)\n",
        "\n",
        "var3 = tf.Variable(tf.linspace(0,1,10))\n",
        "print('\\nvar3: ',var3)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "var1:  <tf.Variable 'Variable:0' shape=() dtype=int32, numpy=4>\n",
            "\n",
            "var2:  <tf.Variable 'Variable:0' shape=(3,) dtype=float32, numpy=array([1., 2., 3.], dtype=float32)>\n",
            "\n",
            "var3:  <tf.Variable 'Variable:0' shape=(10,) dtype=float64, numpy=\n",
            "array([0.        , 0.11111111, 0.22222222, 0.33333333, 0.44444444,\n",
            "       0.55555556, 0.66666667, 0.77777778, 0.88888889, 1.        ])>\n"
          ]
        }
      ],
      "metadata": {
        "id": "CRh-la_X131p",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "57fdee38-987f-4ebf-a427-8a0f6b30aba3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### After construction, the type and shape of the variable are fixed. The value can be changed using one of the assign methods."
      ],
      "metadata": {
        "id": "KV6ZxudhAOZe"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "source": [
        "var4 = tf.Variable([1.0,2.0])\n",
        "print('\\nvar4: ', var4.numpy())\n",
        "\n",
        "var4.assign([3.0,4.0])\n",
        "print('\\nmodified var4: ',var4.numpy())\n",
        "\n",
        "var4.assign_add([1.0,2.0])\n",
        "print('\\nmodified var4: ',var4.numpy())\n",
        "\n",
        "var4.assign_sub([1.0,2.0])\n",
        "print('\\nmodified var4: ',var4.numpy())\n",
        "\n",
        "print('\\nrecall that the shape is fixed')\n",
        "var4.assign([3.0,4.0,4.0])"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "var4:  [1. 2.]\n",
            "\n",
            "modified var4:  [3. 4.]\n",
            "\n",
            "modified var4:  [4. 6.]\n",
            "\n",
            "modified var4:  [3. 4.]\n",
            "\n",
            "recall that the shape is fixed\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "Cannot assign to variable Variable:0 due to variable shape (2,) and value shape (3,) are incompatible",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m/var/folders/6h/mmvp5df90fb3fsqtc5d9mxt40000gn/T/ipykernel_1503/3591645959.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'\\nrecall that the shape is fixed'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0mvar4\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0massign\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m3.0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m4.0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m4.0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m~/miniforge3/envs/tf/lib/python3.8/site-packages/tensorflow/python/ops/resource_variable_ops.py\u001b[0m in \u001b[0;36massign\u001b[0;34m(self, value, use_locking, name, read_value)\u001b[0m\n\u001b[1;32m    896\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    897\u001b[0m           \u001b[0mtensor_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\" \"\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 898\u001b[0;31m         raise ValueError(\n\u001b[0m\u001b[1;32m    899\u001b[0m             (\"Cannot assign to variable%s due to variable shape %s and value \"\n\u001b[1;32m    900\u001b[0m              \"shape %s are incompatible\") %\n",
            "\u001b[0;31mValueError\u001b[0m: Cannot assign to variable Variable:0 due to variable shape (2,) and value shape (3,) are incompatible"
          ]
        }
      ],
      "metadata": {
        "id": "2efeQVVBANxl",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 496
        },
        "outputId": "0965655a-8cb5-436c-bfe1-e8d973c13f20"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gradients in Tensorflow:\n",
        "\n",
        "### To differentiate automatically, TensorFlow needs to remember what operations happen in what order during the forward pass. Then, during the backward pass, TensorFlow traverses this list of operations in reverse order to compute gradients.\n",
        "\n",
        "### We use the tf.GradientTape api for this\n",
        "\n",
        "#### (see detailed guide here https://www.tensorflow.org/guide/autodiff)"
      ],
      "metadata": {
        "id": "N_YfrXkoVTdG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Differentiating wrt a scalar variable"
      ],
      "metadata": {
        "id": "gkK-xns8VTdI"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "source": [
        "x = tf.Variable(4.0)\n",
        "print(' x:',x)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "    y = x**2\n",
        "    z = 2*y\n",
        "\n",
        "print(f'\\n y:', y) # NOTE THE OBJECT TYPE\n",
        "print(f'\\n z:', z)\n",
        "\n",
        "dz_dx = tape.gradient(z, x)  # Evaluate gradient\n",
        "print(f'\\n dz_dx:',dz_dx.numpy())"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " x: <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=4.0>\n",
            "\n",
            " y: tf.Tensor(16.0, shape=(), dtype=float32)\n",
            "\n",
            " z: tf.Tensor(32.0, shape=(), dtype=float32)\n",
            "\n",
            " dz_dx: 16.0\n"
          ]
        }
      ],
      "metadata": {
        "id": "eGdt58rgVTdJ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b548f287-5bd6-446f-bf69-a644f70b0f3d"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Differentiating wrt a scalar tensor (not a variable)"
      ],
      "metadata": {
        "id": "LIUbIw9nVTdM"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.constant(4.0)\n",
        "print(' x:',x)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "    y = x**2\n",
        "    z = 2*y\n",
        "\n",
        "dz_dx = tape.gradient(z, x)  # Evaluate gradient\n",
        "print(f'\\n dz_dx:',dz_dx)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " x: tf.Tensor(4.0, shape=(), dtype=float32)\n",
            "\n",
            " dz_dx: None\n"
          ]
        }
      ],
      "metadata": {
        "id": "jNHRZ0zoVTdM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "234764d3-ea59-4a35-f547-03c31b541a7d"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### To differentiate wrt a tensor x, it needs to be watched by the tape"
      ],
      "metadata": {
        "id": "238giV5zVTdN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.constant(4.0)\n",
        "print(' x:',x)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "    tape.watch(x)\n",
        "    y = x**2\n",
        "    z = 2*y\n",
        "\n",
        "dz_dx = tape.gradient(z, x)  # Evaluate gradient\n",
        "print(f'\\n dz_dx:',dz_dx)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            " x: tf.Tensor(4.0, shape=(), dtype=float32)\n",
            "\n",
            " dz_dx: tf.Tensor(16.0, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "metadata": {
        "id": "_qvM4Mn8VTdO",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0d4ecf0d-da85-4348-ee2d-2d6d028c3dc0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Non-persistant vs persistent tapes\n",
        "### By default, the tape is deleted after a single call to tape.gradient()"
      ],
      "metadata": {
        "id": "cLCEMwfvVTdO"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable(4.0)\n",
        "with tf.GradientTape() as tape:\n",
        "    y = x**2\n",
        "    z = 2*y\n",
        "\n",
        "dz_dx = tape.gradient(z, x) #(4 * x)\n",
        "print(f'\\n dz_dx:',dz_dx)  \n",
        "dy_dx = tape.gradient(y, x)  #(2 * x)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " dz_dx: tf.Tensor(16.0, shape=(), dtype=float32)\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-10-da1e95eecf61>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mdz_dx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#(4 * x)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'\\n dz_dx:'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdz_dx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0mdy_dx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m#(2 * x)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/backprop.py\u001b[0m in \u001b[0;36mgradient\u001b[0;34m(self, target, sources, output_gradients, unconnected_gradients)\u001b[0m\n\u001b[1;32m   1030\u001b[0m     \"\"\"\n\u001b[1;32m   1031\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_tape\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1032\u001b[0;31m       raise RuntimeError(\"A non-persistent GradientTape can only be used to \"\n\u001b[0m\u001b[1;32m   1033\u001b[0m                          \"compute one set of gradients (or jacobians)\")\n\u001b[1;32m   1034\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_recording\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: A non-persistent GradientTape can only be used to compute one set of gradients (or jacobians)"
          ]
        }
      ],
      "metadata": {
        "id": "ly2VCSo_VTdP",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 357
        },
        "outputId": "0dbebf92-5388-4cfd-c5c3-15a0734dd180"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Option 1: Use a persistant tape"
      ],
      "metadata": {
        "id": "AZ4JXqulVTdP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable(4.0)\n",
        "with tf.GradientTape(persistent=True) as tape:\n",
        "    y = x**2\n",
        "    z = 2*y\n",
        "\n",
        "dz_dx = tape.gradient(z, x) #(4 * x)\n",
        "print(f'\\n dz_dx:',dz_dx)  \n",
        "dy_dx = tape.gradient(y, x)  #(2 * x)\n",
        "print(f'\\n dy_dx:',dy_dx)\n",
        "\n",
        "del tape # Clearing the tape, since it doesn't happen automatically now"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " dz_dx: tf.Tensor(16.0, shape=(), dtype=float32)\n",
            "\n",
            " dy_dx: tf.Tensor(8.0, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "metadata": {
        "id": "6JBMVnc_VTdQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f2f6528f-412e-40eb-d48e-12817c2b03c8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Option 2: use nested tapes"
      ],
      "metadata": {
        "id": "bfdzCRotVTdQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable(4.0)\n",
        "with tf.GradientTape() as tape1:\n",
        "    with tf.GradientTape() as tape2:\n",
        "      y = x**2\n",
        "      z = 2*y\n",
        "    dz_dx = tape2.gradient(z, x)  # Tape 2 gets deleted after this call\n",
        "dy_dx = tape1.gradient(y, x)      # Tape 1 gets deleted after this call\n",
        "print(f'\\n dz_dx:',dz_dx) \n",
        "print(f'\\n dy_dx:',dy_dx)  "
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " dz_dx: tf.Tensor(16.0, shape=(), dtype=float32)\n",
            "\n",
            " dy_dx: tf.Tensor(8.0, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "metadata": {
        "id": "2laqDGLoVTdQ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f885409e-dcf8-4bca-8d98-8d612ac58cdb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### We can use this strategy to evalaute higher-order derivates"
      ],
      "metadata": {
        "id": "dRBnzEcfM_yM"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable(2.0)\n",
        "with tf.GradientTape() as tape1:\n",
        "    with tf.GradientTape() as tape2:\n",
        "      y = x**4\n",
        "    dy_dx = tape2.gradient(y, x)      # 4*x**3\n",
        "d2y_dx2 = tape1.gradient(dy_dx, x)    # 12*x**2\n",
        "print(f'\\n dy_dx:',dy_dx) \n",
        "print(f'\\n d2y_dx2:',d2y_dx2) "
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            " dy_dx: tf.Tensor(32.0, shape=(), dtype=float32)\n",
            "\n",
            " d2y_dx2: tf.Tensor(48.0, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "metadata": {
        "id": "-9HOouSzNKeG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d384e798-744c-4299-8b99-af229350e2d4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Differentiating a tensor wrt to a tensor\n",
        "### Say $ x \\in \\mathbb{R}^d$ and $y \\in \\mathbb{R}^D$, the $g := grad(y,x) \\in \\mathbb{R}^d$ with $g_i = \\sum_{j=1}^D \\frac{\\partial y_j}{\\partial x_i}$\n",
        "\n",
        "### In the following example\n",
        "### $$x = \\begin{bmatrix}x_1\\\\x_2\\\\x_3\\end{bmatrix}= \\begin{bmatrix}2.0\\\\3.0\\\\4.0\\end{bmatrix}, \\quad A = \\begin{bmatrix} 1.0 & 2.0 & 0.0 \\\\ 1.0 & 1.0 & -1.0\\end{bmatrix}, \\quad y = Ax = \\begin{bmatrix}x_1 + 2x_2\\\\x_1 + x_2-x_3\\end{bmatrix} = \\begin{bmatrix}8\\\\1\\end{bmatrix}$$\n",
        "### $$\\frac{d y}{d x} = \\begin{bmatrix}1 & 1\\\\2 & 1 \\\\0 & -1\\end{bmatrix}^\\top \\quad \\implies grad(y,x) = \\begin{bmatrix}2\\\\3 \\\\-1\\end{bmatrix}$$"
      ],
      "metadata": {
        "id": "HEsF9mc5VTdT"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable([[2.0],[3.0],[4.0]])\n",
        "A = tf.constant([[1.0,2.0,0.0],[1.0,1.0,-1.0]])\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "  y = tf.matmul(A,x)\n",
        "\n",
        "print('\\nx: ',x.numpy())\n",
        "print('\\ny: ',y.numpy())\n",
        "\n",
        "grad = tape.gradient(y, x)\n",
        "print('\\ngrad(y,x): ',grad.numpy())"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "x:  [[2.]\n",
            " [3.]\n",
            " [4.]]\n",
            "\n",
            "y:  [[8.]\n",
            " [1.]]\n",
            "\n",
            "dydx:  [[ 2.]\n",
            " [ 3.]\n",
            " [-1.]]\n"
          ]
        }
      ],
      "metadata": {
        "id": "u8syTA6HpzZk",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd82312b-75cb-44c7-b8ba-8bd95eadcdf4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### To revaluate the full Jacobian, we need to use tape.jacobian (https://www.tensorflow.org/guide/advanced_autodiff#jacobians)\n",
        "### Note that since x in the example below is a tensor of size 3x1 and y is a tensor of size 2x1, the Jacobian is a tensor of size 2x1x3x1. "
      ],
      "metadata": {
        "id": "ePTiZ6RJvOnN"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.Variable([[2.0],[3.0],[4.0]])\n",
        "A = tf.constant([[1.0,2.0,0.0],[1.0,1.0,-1.0]])\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "  y = tf.matmul(A,x)\n",
        "\n",
        "print('\\nx: ',x.numpy())\n",
        "print('\\ny: ',y.numpy())\n",
        "\n",
        "jac = tape.jacobian(y, x)\n",
        "print('\\ndydx: ',jac.numpy())\n",
        "print('\\ndydx is a tensor of shape: ',jac.shape)\n",
        "print('\\nsqueezed dydx: ',tf.squeeze(jac).numpy())"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "x:  [[2.]\n",
            " [3.]\n",
            " [4.]]\n",
            "\n",
            "y:  [[8.]\n",
            " [1.]]\n",
            "\n",
            "dydx:  [[[[ 1.]\n",
            "   [ 2.]\n",
            "   [ 0.]]]\n",
            "\n",
            "\n",
            " [[[ 1.]\n",
            "   [ 1.]\n",
            "   [-1.]]]]\n",
            "\n",
            "dydx is a tensor of shape:  (2, 1, 3, 1)\n",
            "\n",
            "squeezed dydx:  [[ 1.  2.  0.]\n",
            " [ 1.  1. -1.]]\n"
          ]
        }
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lDvKdUMFz5jd",
        "outputId": "15610d34-44c3-4f70-c951-3ba030b604d4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### For an element-wise calculation, the gradient of the sum gives the derivative of each element with respect to its input-element, since each element is independent. Another way to say this is that the Jacobian is a diagonal matrix."
      ],
      "metadata": {
        "id": "-iJa83Yh2U7C"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.linspace(0.0, 4*np.pi, 100)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "    tape.watch(x)\n",
        "    y = tf.math.sin(x)\n",
        "\n",
        "dy_dx = tape.gradient(y,x)\n",
        "\n",
        "plt.plot(x, y, label='y')\n",
        "plt.plot(x, dy_dx, label='dy/dx')\n",
        "plt.legend(loc='lower left',fontsize=15);\n",
        "plt.xlabel('x',fontsize=15);"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 32600 characters]",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "metadata": {
        "id": "ls6qAuzLVTdT",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 284
        },
        "outputId": "13aec7ef-6bea-4955-dac6-7270da019cf7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gradients of ReLU"
      ],
      "metadata": {
        "id": "dTDvLkRc7X-1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "x = tf.linspace(-1,1, 101)\n",
        "\n",
        "with tf.GradientTape() as tape1:\n",
        "    tape1.watch(x)\n",
        "    with tf.GradientTape() as tape2:\n",
        "        tape2.watch(x)\n",
        "        y = tf.nn.relu(x)\n",
        "    dy_dx = tape2.gradient(y,x)\n",
        "d2y_dx2 = tape1.gradient(dy_dx,x)    \n",
        "\n",
        "plt.plot(x, y, label='y')\n",
        "plt.plot(x, dy_dx, label='dy/dx')\n",
        "plt.plot(x, d2y_dx2, label='d2y/dx2')\n",
        "plt.legend(loc='upper left',fontsize=15);\n",
        "plt.xlabel('x',fontsize=15);"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 16168 characters]",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "metadata": {
        "id": "xPAwSAmn61mO",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 284
        },
        "outputId": "44e4a45f-2d05-4cb1-ed7e-0928c1cdc25c"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [],
      "outputs": [],
      "metadata": {
        "id": "sRcZXEj56-IY"
      }
    }
  ]
}