site/en/guide/advanced_autodiff.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFdPvlXBOdUN"
      },
      "source": [
        "# Advanced automatic differentiation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/advanced_autodiff\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/advanced_autodiff.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/advanced_autodiff.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/advanced_autodiff.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8a859404ce7e"
      },
      "source": [
        "The [Introduction to gradients and automatic differentiation](autodiff.ipynb) guide includes everything required to calculate gradients in TensorFlow. This guide focuses on deeper, less common features of the `tf.GradientTape` API."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqR2PQG4ZaZ0"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import matplotlib as mpl\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "mpl.rcParams['figure.figsize'] = (8, 6)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGRJJRi8TCkJ"
      },
      "source": [
        "## Controlling gradient recording\n",
        "\n",
        "In the [automatic differentiation guide](autodiff.ipynb) you saw how to control which variables and tensors are watched by the tape while building the gradient calculation.\n",
        "\n",
        "The tape also has methods to manipulate the recording."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gB_i0VnhQKt2"
      },
      "source": [
        "### Stop recording\n",
        "\n",
        "If you wish to stop recording gradients, you can use `tf.GradientTape.stop_recording` to temporarily suspend recording.\n",
        "\n",
        "This may be useful to reduce overhead if you do not wish to differentiate a complicated operation in the middle of your model.  This could include calculating a metric or an intermediate result:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mhFSYf7uQWxR"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  x_sq = x * x\n",
        "  with t.stop_recording():\n",
        "    y_sq = y * y\n",
        "  z = x_sq + y_sq\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DEHbEZ1h4p8A"
      },
      "source": [
        "### Reset/start recording from scratch\n",
        "\n",
        "If you wish to start over entirely, use `tf.GradientTape.reset`. Simply exiting the gradient tape block and restarting is usually easier to read, but you can use the `reset` method when exiting the tape block is difficult or impossible."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lsMHsmrh4pqM"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "reset = True\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  y_sq = y * y\n",
        "  if reset:\n",
        "    # Throw out all the tape recorded so far.\n",
        "    t.reset()\n",
        "  z = x * x + y_sq\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6zS7cLmS6zMf"
      },
      "source": [
        "## Stop gradient flow with precision\n",
        "\n",
        "In contrast to the global tape controls above, the `tf.stop_gradient` function is much more precise. It can be used to stop gradients from flowing along a particular path, without needing access to the tape itself:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "30qnZMe48BkB"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  y_sq = y**2\n",
        "  z = x**2 + tf.stop_gradient(y_sq)\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mbb-9lnGVngH"
      },
      "source": [
        "## Custom gradients\n",
        "\n",
        "In some cases, you may want to control exactly how gradients are calculated rather than using the default. These situations include:\n",
        "\n",
        "1. There is no defined gradient for a new op you are writing.\n",
        "2. The default calculations are numerically unstable.\n",
        "3. You wish to cache an expensive computation from the forward pass.\n",
        "4. You want to modify  a value (for example, using `tf.clip_by_value` or `tf.math.round`) without modifying the gradient.\n",
        "\n",
        "For the first case, to write a new op you can use `tf.RegisterGradient` to set up your own (refer to the API docs for details). (Note that the gradient registry is global, so change it with caution.)\n",
        "\n",
        "For the latter three cases, you can use `tf.custom_gradient`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oHr31kc_irF_"
      },
      "source": [
        "Here is an example that applies `tf.clip_by_norm` to the intermediate gradient:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mjj01w4NYtwd"
      },
      "outputs": [],
      "source": [
        "# Establish an identity operation, but clip during the gradient pass.\n",
        "@tf.custom_gradient\n",
        "def clip_gradients(y):\n",
        "  def backward(dy):\n",
        "    return tf.clip_by_norm(dy, 0.5)\n",
        "  return y, backward\n",
        "\n",
        "v = tf.Variable(2.0)\n",
        "with tf.GradientTape() as t:\n",
        "  output = clip_gradients(v * v)\n",
        "print(t.gradient(output, v))  # calls \"backward\", which clips 4 to 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4t7S0scYrD3"
      },
      "source": [
        "Refer to the `tf.custom_gradient` decorator API docs for more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v0ODp4Oi--I0"
      },
      "source": [
        "### Custom gradients in SavedModel\n",
        "\n",
        "Note: This feature is available from TensorFlow 2.6.\n",
        "\n",
        "Custom gradients can be saved to SavedModel by using the option `tf.saved_model.SaveOptions(experimental_custom_gradients=True)`.\n",
        "\n",
        "To be saved into the SavedModel, the gradient function must be traceable (to learn more, check out the [Better performance with tf.function](function.ipynb) guide)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q5JBgIBYjN1I"
      },
      "outputs": [],
      "source": [
        "class MyModule(tf.Module):\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec(None)])\n",
        "  def call_custom_grad(self, x):\n",
        "    return clip_gradients(x)\n",
        "\n",
        "model = MyModule()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZTrgy2q-9pq"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(\n",
        "    model,\n",
        "    'saved_model',\n",
        "    options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))\n",
        "\n",
        "# The loaded gradients will be the same as the above example.\n",
        "v = tf.Variable(2.0)\n",
        "loaded = tf.saved_model.load('saved_model')\n",
        "with tf.GradientTape() as t:\n",
        "  output = loaded.call_custom_grad(v * v)\n",
        "print(t.gradient(output, v))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-LfRs5FbJCk"
      },
      "source": [
        "A note about the above example: If you try replacing the above code with `tf.saved_model.SaveOptions(experimental_custom_gradients=False)`, the gradient will still produce the same result on loading. The reason is that the gradient registry still contains the custom gradient used in the function `call_custom_op`. However, if you restart the runtime after saving without custom gradients, running the loaded model under the `tf.GradientTape` will throw the error: `LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN)`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aENEt6Veryb"
      },
      "source": [
        "## Multiple tapes\n",
        "\n",
        "Multiple tapes interact seamlessly.\n",
        "\n",
        "For example, here each tape watches a different set of tensors:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BJ0HdMvte0VZ"
      },
      "outputs": [],
      "source": [
        "x0 = tf.constant(0.0)\n",
        "x1 = tf.constant(0.0)\n",
        "\n",
        "with tf.GradientTape() as tape0, tf.GradientTape() as tape1:\n",
        "  tape0.watch(x0)\n",
        "  tape1.watch(x1)\n",
        "\n",
        "  y0 = tf.math.sin(x0)\n",
        "  y1 = tf.nn.sigmoid(x1)\n",
        "\n",
        "  y = y0 + y1\n",
        "\n",
        "  ys = tf.reduce_sum(y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ApAoMNFfNz6"
      },
      "outputs": [],
      "source": [
        "tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rF1jrAJsfYW_"
      },
      "outputs": [],
      "source": [
        "tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DK05KXrAAld3"
      },
      "source": [
        "### Higher-order gradients\n",
        "\n",
        "Operations inside of the `tf.GradientTape` context manager are recorded for automatic differentiation. If gradients are computed in that context, then the gradient computation is recorded as well. As a result, the exact same API works for higher-order gradients as well.\n",
        "\n",
        "For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cPQgthZ7ugRJ"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(1.0)  # Create a Tensorflow variable initialized to 1.0\n",
        "\n",
        "with tf.GradientTape() as t2:\n",
        "  with tf.GradientTape() as t1:\n",
        "    y = x * x * x\n",
        "\n",
        "  # Compute the gradient inside the outer `t2` context manager\n",
        "  # which means the gradient computation is differentiable as well.\n",
        "  dy_dx = t1.gradient(y, x)\n",
        "d2y_dx2 = t2.gradient(dy_dx, x)\n",
        "\n",
        "print('dy_dx:', dy_dx.numpy())  # 3 * x**2 => 3.0\n",
        "print('d2y_dx2:', d2y_dx2.numpy())  # 6 * x => 6.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k0HV-Ah4_76i"
      },
      "source": [
        "While that does give you the second derivative of a _scalar_ function, this pattern does not generalize to produce a Hessian matrix, since `tf.GradientTape.gradient` only computes the gradient of a scalar. To construct a [Hessian matrix](https://en.wikipedia.org/wiki/Hessian_matrix), go to the [Hessian example](#hessian) under the [Jacobian section](#jacobians).\n",
        "\n",
        "\"Nested calls to `tf.GradientTape.gradient`\" is a good pattern when you are calculating a scalar from a gradient, and then the resulting scalar acts as a source for a second gradient calculation, as in the following example.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t7LRlcpVKHv1"
      },
      "source": [
        "#### Example: Input gradient regularization\n",
        "\n",
        "Many models are susceptible to \"adversarial examples\". This collection of techniques modifies the model's input to confuse the model's output. The simplest implementation—such as the [Adversarial example using the Fast Gradient Signed Method attack](https://www.tensorflow.org/tutorials/generative/adversarial_fgsm)—takes a single step along the gradient of the output with respect to the input; the \"input gradient\".\n",
        "\n",
        "One technique to increase robustness to adversarial examples is [input gradient regularization](https://arxiv.org/abs/1905.11468) (Finlay & Oberman, 2019), which attempts to minimize the magnitude of the input gradient. If the input gradient is small, then the change in the output should be small too.\n",
        "\n",
        "Below is a naive implementation of input gradient regularization. The implementation is:\n",
        "\n",
        "1. Calculate the gradient of the output with respect to the input using an inner tape.\n",
        "2. Calculate the magnitude of that input gradient.\n",
        "3. Calculate the gradient of that magnitude with respect to the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tH3ZFuUfDLrR"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E6yOFsjEDR9u"
      },
      "outputs": [],
      "source": [
        "with tf.GradientTape() as t2:\n",
        "  # The inner tape only takes the gradient with respect to the input,\n",
        "  # not the variables.\n",
        "  with tf.GradientTape(watch_accessed_variables=False) as t1:\n",
        "    t1.watch(x)\n",
        "    y = layer(x)\n",
        "    out = tf.reduce_sum(layer(x)**2)\n",
        "  # 1. Calculate the input gradient.\n",
        "  g1 = t1.gradient(out, x)\n",
        "  # 2. Calculate the magnitude of the input gradient.\n",
        "  g1_mag = tf.norm(g1)\n",
        "\n",
        "# 3. Calculate the gradient of the magnitude with respect to the model.\n",
        "dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "123QMq6PqK_d"
      },
      "outputs": [],
      "source": [
        "[var.shape for var in dg1_mag]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4xiYigexMtQ"
      },
      "source": [
        "## Jacobians\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4-hVHVIeExkI"
      },
      "source": [
        "All the previous examples took the gradients of a scalar target with respect to some source tensor(s).\n",
        "\n",
        "The [Jacobian matrix](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) represents the gradients of a vector valued function. Each row contains the gradient of one of the vector's elements.\n",
        "\n",
        "The `tf.GradientTape.jacobian` method allows you to efficiently calculate a Jacobian matrix."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KzNyIM0QBYIH"
      },
      "source": [
        "Note that:\n",
        "\n",
        "* Like `gradient`: The `sources` argument can be a tensor or a container of tensors.\n",
        "* Unlike `gradient`: The `target` tensor must be a single tensor."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O74K3hlxBC8a"
      },
      "source": [
        "### Scalar source"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B08OKn1Orkuc"
      },
      "source": [
        "As a first example, here is the Jacobian of a vector-target with respect to a scalar-source."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bAFeIE8EuVIq"
      },
      "outputs": [],
      "source": [
        "x = tf.linspace(-10.0, 10.0, 200+1)\n",
        "delta = tf.Variable(0.0)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "  y = tf.nn.sigmoid(x+delta)\n",
        "\n",
        "dy_dx = tape.jacobian(y, delta)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgHbUk3zr-WU"
      },
      "source": [
        "When you take the Jacobian with respect to a scalar the result has the shape of the **target**, and gives the gradient of the each element with respect to the source:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZ6awnDzr_BA"
      },
      "outputs": [],
      "source": [
        "print(y.shape)\n",
        "print(dy_dx.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "siNZaklc0_-e"
      },
      "outputs": [],
      "source": [
        "plt.plot(x.numpy(), y, label='y')\n",
        "plt.plot(x.numpy(), dy_dx, label='dy/dx')\n",
        "plt.legend()\n",
        "_ = plt.xlabel('x')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsOMSD_1BGkD"
      },
      "source": [
        "### Tensor source"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g3iXKN7KF-st"
      },
      "source": [
        "Whether the input is scalar or tensor, `tf.GradientTape.jacobian` efficiently calculates the gradient of each element of the source with respect to each element of the target(s).\n",
        "\n",
        "For example, the output of this layer has a shape of `(10, 7)`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39YXItgLxMBk"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)\n",
        "\n",
        "with tf.GradientTape(persistent=True) as tape:\n",
        "  y = layer(x)\n",
        "\n",
        "y.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tshNRtfKuVP_"
      },
      "source": [
        "And the layer's kernel's shape is `(5, 10)`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CigTWyfPvPuv"
      },
      "outputs": [],
      "source": [
        "layer.kernel.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mN96JRpnAjpx"
      },
      "source": [
        "The shape of the Jacobian of the output with respect to the kernel is those two shapes concatenated together:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRLzTTbvEimH"
      },
      "outputs": [],
      "source": [
        "j = tape.jacobian(y, layer.kernel)\n",
        "j.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Lrv7miMvTll"
      },
      "source": [
        "If you sum over the target's dimensions, you're left with the gradient of the sum that would have been calculated by `tf.GradientTape.gradient`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FJjZpYRnDjVa"
      },
      "outputs": [],
      "source": [
        "g = tape.gradient(y, layer.kernel)\n",
        "print('g.shape:', g.shape)\n",
        "\n",
        "j_sum = tf.reduce_sum(j, axis=[0, 1])\n",
        "delta = tf.reduce_max(abs(g - j_sum)).numpy()\n",
        "assert delta < 1e-3\n",
        "print('delta:', delta)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZKajuGlk_krs"
      },
      "source": [
        "<a id=\"hessian\"> </hessian>\n",
        "\n",
        "#### Example: Hessian"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NYcsXeo8TDLi"
      },
      "source": [
        "While `tf.GradientTape` doesn't give an explicit method for constructing a [Hessian matrix](https://en.wikipedia.org/wiki/Hessian_matrix) it's possible to build one using the `tf.GradientTape.jacobian` method.\n",
        "\n",
        "Note: The Hessian matrix contains `N**2` parameters. For this and other reasons it is not practical for most models. This example is included more as a demonstration of how to use the `tf.GradientTape.jacobian` method, and is not an endorsement of direct Hessian-based optimization. A Hessian-vector product can be [calculated efficiently with nested tapes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py), and is a much more efficient approach to second-order optimization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ELGTaell_j81"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.relu)\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.relu)\n",
        "\n",
        "with tf.GradientTape() as t2:\n",
        "  with tf.GradientTape() as t1:\n",
        "    x = layer1(x)\n",
        "    x = layer2(x)\n",
        "    loss = tf.reduce_mean(x**2)\n",
        "\n",
        "  g = t1.gradient(loss, layer1.kernel)\n",
        "\n",
        "h = t2.jacobian(g, layer1.kernel)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FVqQuZj4XGjm"
      },
      "outputs": [],
      "source": [
        "print(f'layer.kernel.shape: {layer1.kernel.shape}')\n",
        "print(f'h.shape: {h.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_M7XElgaiMeP"
      },
      "source": [
        "To use this Hessian for a [Newton's method](https://en.wikipedia.org/wiki/Newton%27s_method_in_optimization) step, you would first flatten out its axes into a matrix, and flatten out the gradient into a vector:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6te7N6wVXwXX"
      },
      "outputs": [],
      "source": [
        "n_params = tf.reduce_prod(layer1.kernel.shape)\n",
        "\n",
        "g_vec = tf.reshape(g, [n_params, 1])\n",
        "h_mat = tf.reshape(h, [n_params, n_params])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9rO8b-0mgOH"
      },
      "source": [
        "The Hessian matrix should be symmetric:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8TCHc7Vrf52S"
      },
      "outputs": [],
      "source": [
        "def imshow_zero_center(image, **kwargs):\n",
        "  lim = tf.reduce_max(abs(image))\n",
        "  plt.imshow(image, vmin=-lim, vmax=lim, cmap='seismic', **kwargs)\n",
        "  plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DExOxd7Ok2H0"
      },
      "outputs": [],
      "source": [
        "imshow_zero_center(h_mat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "13fBswmtQes4"
      },
      "source": [
        "The Newton's method update step is shown below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3DdnbynBdSor"
      },
      "outputs": [],
      "source": [
        "eps = 1e-3\n",
        "eye_eps = tf.eye(h_mat.shape[0])*eps"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-zPdtyoWeUeV"
      },
      "source": [
        "Note: [Don't  actually invert the matrix](https://www.johndcook.com/blog/2010/01/19/dont-invert-that-matrix/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k1LYftgmswOO"
      },
      "outputs": [],
      "source": [
        "# X(k+1) = X(k) - (∇²f(X(k)))^-1 @ ∇f(X(k))\n",
        "# h_mat = ∇²f(X(k))\n",
        "# g_vec = ∇f(X(k))\n",
        "update = tf.linalg.solve(h_mat + eye_eps, g_vec)\n",
        "\n",
        "# Reshape the update and apply it to the variable.\n",
        "_ = layer1.kernel.assign_sub(tf.reshape(update, layer1.kernel.shape))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pF6qjlHKWxF4"
      },
      "source": [
        "While this is relatively simple for a single `tf.Variable`, applying this to a non-trivial model would require careful concatenation and slicing to produce a full Hessian across multiple variables."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PQWM0uN-GO5t"
      },
      "source": [
        "### Batch Jacobian"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hKtB3rY6EySJ"
      },
      "source": [
        "In some cases, you want to take the Jacobian of each of a stack of targets with respect to a stack of sources, where the Jacobians for each target-source pair are independent.\n",
        "\n",
        "For example, here the input `x` is shaped `(batch, ins)` and the output `y` is shaped `(batch, outs)`:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tQMndhIUHMes"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)\n",
        "\n",
        "with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:\n",
        "  tape.watch(x)\n",
        "  y = layer1(x)\n",
        "  y = layer2(y)\n",
        "\n",
        "y.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ff2spRHEJXBU"
      },
      "source": [
        "The full Jacobian of `y` with respect to `x` has a shape of `(batch, ins, batch, outs)`, even if you only want `(batch, ins, outs)`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1zSl2A5-HhMH"
      },
      "outputs": [],
      "source": [
        "j = tape.jacobian(y, x)\n",
        "j.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UibJijPLJrpQ"
      },
      "source": [
        "If the gradients of each item in the stack are independent, then every `(batch, batch)` slice of this tensor is a diagonal matrix:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFl9uj3ueVSH"
      },
      "outputs": [],
      "source": [
        "imshow_zero_center(j[:, 0, :, 0])\n",
        "_ = plt.title('A (batch, batch) slice')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g4ZoRJcJNmy5"
      },
      "outputs": [],
      "source": [
        "def plot_as_patches(j):\n",
        "  # Reorder axes so the diagonals will each form a contiguous patch.\n",
        "  j = tf.transpose(j, [1, 0, 3, 2])\n",
        "  # Pad in between each patch.\n",
        "  lim = tf.reduce_max(abs(j))\n",
        "  j = tf.pad(j, [[0, 0], [1, 1], [0, 0], [1, 1]],\n",
        "             constant_values=-lim)\n",
        "  # Reshape to form a single image.\n",
        "  s = j.shape\n",
        "  j = tf.reshape(j, [s[0]*s[1], s[2]*s[3]])\n",
        "  imshow_zero_center(j, extent=[-0.5, s[2]-0.5, s[0]-0.5, -0.5])\n",
        "\n",
        "plot_as_patches(j)\n",
        "_ = plt.title('All (batch, batch) slices are diagonal')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OXpTBKyeK84z"
      },
      "source": [
        "To get the desired result, you can sum over the duplicate `batch` dimension, or else select the diagonals using `tf.einsum`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v65OAjEgLQwl"
      },
      "outputs": [],
      "source": [
        "j_sum = tf.reduce_sum(j, axis=2)\n",
        "print(j_sum.shape)\n",
        "j_select = tf.einsum('bxby->bxy', j)\n",
        "print(j_select.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zT_VfR6lcwxD"
      },
      "source": [
        "It would be much more efficient to do the calculation without the extra dimension in the first place. The `tf.GradientTape.batch_jacobian` method does exactly that:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YJLIl9WpHqYq"
      },
      "outputs": [],
      "source": [
        "jb = tape.batch_jacobian(y, x)\n",
        "jb.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-5t_q5SfHw7T"
      },
      "outputs": [],
      "source": [
        "error = tf.reduce_max(abs(jb - j_sum))\n",
        "assert error < 1e-3\n",
        "print(error.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IUeY2ZCiL31I"
      },
      "source": [
        "Caution: `tf.GradientTape.batch_jacobian` only verifies that the first dimension of the source and target match. It doesn't check that the gradients are actually independent. It's up to you to make sure you only use `batch_jacobian` where it makes sense. For example, adding a `tf.keras.layers.BatchNormalization` destroys the independence, since it normalizes across the `batch` dimension:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnDugVc-L4fj"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)\n",
        "bn = tf.keras.layers.BatchNormalization()\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)\n",
        "\n",
        "with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:\n",
        "  tape.watch(x)\n",
        "  y = layer1(x)\n",
        "  y = bn(y, training=True)\n",
        "  y = layer2(y)\n",
        "\n",
        "j = tape.jacobian(y, x)\n",
        "print(f'j.shape: {j.shape}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SNyZ1WhJMVLm"
      },
      "outputs": [],
      "source": [
        "plot_as_patches(j)\n",
        "\n",
        "_ = plt.title('These slices are not diagonal')\n",
        "_ = plt.xlabel(\"Don't use `batch_jacobian`\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_x7ih5sarvG"
      },
      "source": [
        "In this case, `batch_jacobian` still runs and returns _something_ with the expected shape, but its contents have an unclear meaning:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k8_mICHoasCi"
      },
      "outputs": [],
      "source": [
        "jb = tape.batch_jacobian(y, x)\n",
        "print(f'jb.shape: {jb.shape}')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "advanced_autodiff.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}