site/en/guide/intro_to_graphs.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N7ITxKLUkX0v"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yOYx6tzSnWQ3"
      },
      "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": "6xgB0Oz5eGSQ"
      },
      "source": [
        "# Introduction to graphs and tf.function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4zzZVZtQb1w"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/intro_to_graphs\"><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/intro_to_graphs.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/intro_to_graphs.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/intro_to_graphs.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RBKqnXI9GOax"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This guide goes beneath the surface of TensorFlow and Keras to demonstrate how TensorFlow works. If you instead want to immediately get started with Keras, check out the [collection of Keras guides](https://www.tensorflow.org/guide/keras/).\n",
        "\n",
        "In this guide, you'll learn how TensorFlow allows you to make simple changes to your code to get graphs, how graphs are stored and represented, and how you can use them to accelerate your models.\n",
        "\n",
        "Note: For those of you who are only familiar with TensorFlow 1.x, this guide demonstrates a very different view of graphs.\n",
        "\n",
        "**This is a big-picture overview that covers how `tf.function` allows you to switch from eager execution to graph execution.** For a more complete specification of `tf.function`, go to the [Better performance with `tf.function`](./function.ipynb) guide.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v0DdlfacAdTZ"
      },
      "source": [
        "### What are graphs?\n",
        "\n",
        "In the previous three guides, you ran TensorFlow **eagerly**. This means TensorFlow operations are executed by Python, operation by operation, and return results back to Python.\n",
        "\n",
        "While eager execution has several unique advantages, graph execution enables portability outside Python and tends to offer better performance. **Graph execution** means that tensor computations are executed as a *TensorFlow graph*, sometimes referred to as a `tf.Graph` or simply a \"graph.\"\n",
        "\n",
        "**Graphs are data structures that contain a set of `tf.Operation` objects, which represent units of computation; and `tf.Tensor` objects, which represent the units of data that flow between operations.** They are defined in a `tf.Graph` context. Since these graphs are data structures, they can be saved, run, and restored all without the original Python code.\n",
        "\n",
        "This is what a TensorFlow graph representing a two-layer neural network looks like when visualized in TensorBoard:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FvQ5aBuRGT1o"
      },
      "source": [
        "<img alt=\"A simple TensorFlow graph\" src=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/images/intro_to_graphs/two-layer-network.png?raw=1\">"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHpY3avXGITP"
      },
      "source": [
        "### The benefits of graphs\n",
        "\n",
        "With a graph, you have a great deal of flexibility.  You can use your TensorFlow graph in environments that don't have a Python interpreter, like mobile applications, embedded devices, and backend servers. TensorFlow uses graphs as the format for [saved models](./saved_model.ipynb) when it exports them from Python.\n",
        "\n",
        "Graphs are also easily optimized, allowing the compiler to do transformations like:\n",
        "\n",
        "* Statically infer the value of tensors by folding constant nodes in your computation *(\"constant folding\")*.\n",
        "* Separate sub-parts of a computation that are independent and split them between threads or devices.\n",
        "* Simplify arithmetic operations by eliminating common subexpressions.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o1x1EOD9GjnB"
      },
      "source": [
        "There is an entire optimization system, [Grappler](./graph_optimization.ipynb), to perform this and other speedups.\n",
        "\n",
        "In short, graphs are extremely useful and let your TensorFlow run **fast**, run **in parallel**, and run efficiently **on multiple devices**.\n",
        "\n",
        "However, you still want to define your machine learning models (or other computations) in Python for convenience, and then automatically construct graphs when you need them."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k-6Qi0thw2i9"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d1689fa928f"
      },
      "source": [
        "Import some necessary libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "goZwOXp_xyQj"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import timeit\n",
        "from datetime import datetime"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pSZebVuWxDXu"
      },
      "source": [
        "## Taking advantage of graphs\n",
        "\n",
        "You create and run a graph in TensorFlow by using `tf.function`, either as a direct call or as a decorator. `tf.function` takes a regular function as input and returns a `tf.types.experimental.PolymorphicFunction`. **A `PolymorphicFunction` is a Python callable that builds TensorFlow graphs from the Python function. You use a `tf.function` in the same way as its Python equivalent.**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HKbLeJ1y0Umi"
      },
      "outputs": [],
      "source": [
        "# Define a Python function.\n",
        "def a_regular_function(x, y, b):\n",
        "  x = tf.matmul(x, y)\n",
        "  x = x + b\n",
        "  return x\n",
        "\n",
        "# The Python type of `a_function_that_uses_a_graph` will now be a\n",
        "# `PolymorphicFunction`.\n",
        "a_function_that_uses_a_graph = tf.function(a_regular_function)\n",
        "\n",
        "# Make some tensors.\n",
        "x1 = tf.constant([[1.0, 2.0]])\n",
        "y1 = tf.constant([[2.0], [3.0]])\n",
        "b1 = tf.constant(4.0)\n",
        "\n",
        "orig_value = a_regular_function(x1, y1, b1).numpy()\n",
        "# Call a `tf.function` like a Python function.\n",
        "tf_function_value = a_function_that_uses_a_graph(x1, y1, b1).numpy()\n",
        "assert(orig_value == tf_function_value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNvuAYpdrTOf"
      },
      "source": [
        "On the outside, a `tf.function` looks like a regular function you write using TensorFlow operations. [Underneath](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/eager/polymorphic_function/polymorphic_function.py), however, it is *very different*. The underlying `PolymorphicFunction` **encapsulates several `tf.Graph`s behind one API** (learn more in the _Polymorphism_ section). That is how a `tf.function` is able to give you the benefits of graph execution, like speed and deployability (refer to _The benefits of graphs_ above)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MT7U8ozok0gV"
      },
      "source": [
        "`tf.function` applies to a function *and all other functions it calls*:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rpz08iLplm9F"
      },
      "outputs": [],
      "source": [
        "def inner_function(x, y, b):\n",
        "  x = tf.matmul(x, y)\n",
        "  x = x + b\n",
        "  return x\n",
        "\n",
        "# Using the `tf.function` decorator makes `outer_function` into a\n",
        "# `PolymorphicFunction`.\n",
        "@tf.function\n",
        "def outer_function(x):\n",
        "  y = tf.constant([[2.0], [3.0]])\n",
        "  b = tf.constant(4.0)\n",
        "\n",
        "  return inner_function(x, y, b)\n",
        "\n",
        "# Note that the callable will create a graph that\n",
        "# includes `inner_function` as well as `outer_function`.\n",
        "outer_function(tf.constant([[1.0, 2.0]])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P88fOr88qgCj"
      },
      "source": [
        "If you have used TensorFlow 1.x, you will notice that at no time did you need to define a `Placeholder` or `tf.Session`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wfeKf0Nr1OEK"
      },
      "source": [
        "### Converting Python functions to graphs\n",
        "\n",
        "Any function you write with TensorFlow will contain a mixture of built-in TF operations and Python logic, such as `if-then` clauses, loops, `break`, `return`, `continue`, and more. While TensorFlow operations are easily captured by a `tf.Graph`, Python-specific logic needs to undergo an extra step in order to become part of the graph. `tf.function` uses a library called AutoGraph (`tf.autograph`) to convert Python code into graph-generating code.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PFObpff1BMEb"
      },
      "outputs": [],
      "source": [
        "def simple_relu(x):\n",
        "  if tf.greater(x, 0):\n",
        "    return x\n",
        "  else:\n",
        "    return 0\n",
        "\n",
        "# Using `tf.function` makes `tf_simple_relu` a `PolymorphicFunction` that wraps\n",
        "# `simple_relu`.\n",
        "tf_simple_relu = tf.function(simple_relu)\n",
        "\n",
        "print(\"First branch, with graph:\", tf_simple_relu(tf.constant(1)).numpy())\n",
        "print(\"Second branch, with graph:\", tf_simple_relu(tf.constant(-1)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hO4DBUNZBMwQ"
      },
      "source": [
        "Though it is unlikely that you will need to view graphs directly, you can inspect the outputs to check the exact results. These are not easy to read, so no need to look too carefully!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lAKaat3w0gnn"
      },
      "outputs": [],
      "source": [
        "# This is the graph-generating output of AutoGraph.\n",
        "print(tf.autograph.to_code(simple_relu))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8x6RAqza1UWf"
      },
      "outputs": [],
      "source": [
        "# This is the graph itself.\n",
        "print(tf_simple_relu.get_concrete_function(tf.constant(1)).graph.as_graph_def())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ4Ieg6tBE6l"
      },
      "source": [
        "Most of the time, `tf.function` will work without  special considerations. However, there are some caveats, and the [`tf.function` guide](./function.ipynb) can help here, as well as the [complete AutoGraph reference](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sIpc_jfjEZEg"
      },
      "source": [
        "### Polymorphism: one `tf.function`, many graphs\n",
        "\n",
        "A `tf.Graph` is specialized to a specific type of inputs (for example, tensors with a specific [`dtype`](https://www.tensorflow.org/api_docs/python/tf/dtypes/DType) or objects with the same [`id()`](https://docs.python.org/3/library/functions.html#id)).\n",
        "\n",
        "Each time you invoke a `tf.function` with a set of arguments that can't be handled by any of its existing graphs (such as arguments with new `dtypes` or incompatible shapes), it creates a new `tf.Graph` specialized to those new arguments. The type specification of a `tf.Graph`'s inputs is represented by `tf.types.experimental.FunctionType`, also referred to as the **signature**. For more information regarding when a new `tf.Graph` is generated, how that can be controlled, and how `FunctionType` can be useful, go to the _Rules of tracing_ section of the [Better performance with `tf.function`](./function.ipynb) guide.\n",
        "\n",
        "The `tf.function` stores the `tf.Graph` corresponding to that signature in a `ConcreteFunction`. **A `ConcreteFunction` can be thought of as a wrapper around a `tf.Graph`.**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LOASwhbvIv_T"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def my_relu(x):\n",
        "  return tf.maximum(0., x)\n",
        "\n",
        "# `my_relu` creates new graphs as it observes different input types.\n",
        "print(my_relu(tf.constant(5.5)))\n",
        "print(my_relu([1, -1]))\n",
        "print(my_relu(tf.constant([3., -3.])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1qRtw7R4KL9X"
      },
      "source": [
        "If the `tf.function` has already been called with the same input types, it does not create a new `tf.Graph`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TjjbnL5OKNDP"
      },
      "outputs": [],
      "source": [
        "# These two calls do *not* create new graphs.\n",
        "print(my_relu(tf.constant(-2.5))) # Input type matches `tf.constant(5.5)`.\n",
        "print(my_relu(tf.constant([-1., 1.]))) # Input type matches `tf.constant([3., -3.])`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UohRmexhIpvQ"
      },
      "source": [
        "Because it's backed by multiple graphs, a `tf.function` is (as the name \"PolymorphicFunction\" suggests) **polymorphic**. That enables it to support more input types than a single `tf.Graph` could represent, and to optimize each `tf.Graph` for better performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dxzqebDYFmLy"
      },
      "outputs": [],
      "source": [
        "# There are three `ConcreteFunction`s (one for each graph) in `my_relu`.\n",
        "# The `ConcreteFunction` also knows the return type and shape!\n",
        "print(my_relu.pretty_printed_concrete_signatures())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V11zkxU22XeD"
      },
      "source": [
        "## Using `tf.function`\n",
        "\n",
        "So far, you've learned how to convert a Python function into a graph simply by using `tf.function` as a decorator or wrapper. But in practice, getting `tf.function` to work correctly can be tricky! In the following sections, you'll learn how you can make your code work as expected with `tf.function`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yp_n0B5-P0RU"
      },
      "source": [
        "### Graph execution vs. eager execution\n",
        "\n",
        "The code in a `tf.function` can be executed both eagerly and as a graph. By default, `tf.function` executes its code as a graph:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_R0BOvBFxqVZ"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def get_MSE(y_true, y_pred):\n",
        "  sq_diff = tf.pow(y_true - y_pred, 2)\n",
        "  return tf.reduce_mean(sq_diff)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zikMVPGhmDET"
      },
      "outputs": [],
      "source": [
        "y_true = tf.random.uniform([5], maxval=10, dtype=tf.int32)\n",
        "y_pred = tf.random.uniform([5], maxval=10, dtype=tf.int32)\n",
        "print(y_true)\n",
        "print(y_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "07r08Dh158ft"
      },
      "outputs": [],
      "source": [
        "get_MSE(y_true, y_pred)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cyZNCRcQorGO"
      },
      "source": [
        "To verify that your `tf.function`'s graph is doing the same computation as its equivalent Python function, you can make it execute eagerly with `tf.config.run_functions_eagerly(True)`. This is a switch that **turns off `tf.function`'s ability to create and run graphs**, instead of executing the code normally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lKoF6NjPoI8w"
      },
      "outputs": [],
      "source": [
        "tf.config.run_functions_eagerly(True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ZLqTyn0oKeM"
      },
      "outputs": [],
      "source": [
        "get_MSE(y_true, y_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cV7daQW9odn-"
      },
      "outputs": [],
      "source": [
        "# Don't forget to set it back when you are done.\n",
        "tf.config.run_functions_eagerly(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DKT3YBsqy0x4"
      },
      "source": [
        "However, `tf.function` can behave differently under graph and eager execution. The Python [`print`](https://docs.python.org/3/library/functions.html#print) function is one example of how these two modes differ. Let's check out what happens when you insert a `print` statement to your function and call it repeatedly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BEJeVeBEoGjV"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def get_MSE(y_true, y_pred):\n",
        "  print(\"Calculating MSE!\")\n",
        "  sq_diff = tf.pow(y_true - y_pred, 2)\n",
        "  return tf.reduce_mean(sq_diff)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3sWTGwX3BzP1"
      },
      "source": [
        "Observe what is printed:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3rJIeBg72T9n"
      },
      "outputs": [],
      "source": [
        "error = get_MSE(y_true, y_pred)\n",
        "error = get_MSE(y_true, y_pred)\n",
        "error = get_MSE(y_true, y_pred)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WLMXk1uxKQ44"
      },
      "source": [
        "Is the output surprising? **`get_MSE` only printed once even though it was called *three* times.**\n",
        "\n",
        "To explain, the `print` statement is executed when `tf.function` runs the original code in order to create the graph in a process known as \"tracing\" (refer to the _Tracing_ section of the [`tf.function` guide](./function.ipynb). **Tracing captures the TensorFlow operations into a graph, and `print` is not captured in the graph.**  That graph is then executed for all three calls **without ever running the Python code again**.\n",
        "\n",
        "As a sanity check, let's turn off graph execution to compare:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oFSxRtcptYpe"
      },
      "outputs": [],
      "source": [
        "# Now, globally set everything to run eagerly to force eager execution.\n",
        "tf.config.run_functions_eagerly(True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qYxrAtvzNgHR"
      },
      "outputs": [],
      "source": [
        "# Observe what is printed below.\n",
        "error = get_MSE(y_true, y_pred)\n",
        "error = get_MSE(y_true, y_pred)\n",
        "error = get_MSE(y_true, y_pred)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Df6ynXcAaup"
      },
      "outputs": [],
      "source": [
        "tf.config.run_functions_eagerly(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PUR7qC_bquCn"
      },
      "source": [
        "`print` is a *Python side effect*, and there are other differences that you should be aware of when converting a function into a `tf.function`. Learn more in the _Limitations_ section of the [Better performance with `tf.function`](./function.ipynb) guide."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oTZJfV_tccVp"
      },
      "source": [
        "Note: If you would like to print values in both eager and graph execution, use `tf.print` instead."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rMT_Xf5yKn9o"
      },
      "source": [
        "### Non-strict execution\n",
        "\n",
        "<a id=\"non-strict\"></a>\n",
        "\n",
        "Graph execution only executes the operations necessary to produce the observable effects, which include:\n",
        "\n",
        "- The return value of the function\n",
        "- Documented well-known side-effects such as:\n",
        "  - Input/output operations, like `tf.print`\n",
        "  - Debugging operations, such as the assert functions in `tf.debugging`\n",
        "  - Mutations of `tf.Variable`\n",
        "\n",
        "This behavior is usually known as \"Non-strict execution\", and differs from eager execution, which steps through all of the program operations, needed or not.\n",
        "\n",
        "In particular, runtime error checking does not count as an observable effect. If an operation is skipped because it is unnecessary, it cannot raise any runtime errors.\n",
        "\n",
        "In the following example, the \"unnecessary\" operation `tf.gather` is skipped during graph execution, so the runtime error `InvalidArgumentError` is not raised as it would be in eager execution. Do not rely on an error being raised while executing a graph."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OdN0nKlUwj7M"
      },
      "outputs": [],
      "source": [
        "def unused_return_eager(x):\n",
        "  # Get index 1 will fail when `len(x) == 1`\n",
        "  tf.gather(x, [1]) # unused \n",
        "  return x\n",
        "\n",
        "try:\n",
        "  print(unused_return_eager(tf.constant([0.0])))\n",
        "except tf.errors.InvalidArgumentError as e:\n",
        "  # All operations are run during eager execution so an error is raised.\n",
        "  print(f'{type(e).__name__}: {e}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d80Fob4MwhTs"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def unused_return_graph(x):\n",
        "  tf.gather(x, [1]) # unused\n",
        "  return x\n",
        "\n",
        "# Only needed operations are run during graph execution. The error is not raised.\n",
        "print(unused_return_graph(tf.constant([0.0])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "def6MupG9R0O"
      },
      "source": [
        "### `tf.function` best practices\n",
        "\n",
        "It may take some time to get used to the behavior of `tf.function`.  To get started quickly, first-time users should play around with decorating toy functions with `@tf.function` to get experience with going from eager to graph execution.\n",
        "\n",
        "*Designing for `tf.function`* may be your best bet for writing graph-compatible TensorFlow programs. Here are some tips:\n",
        "-  Toggle between eager and graph execution early and often with `tf.config.run_functions_eagerly` to pinpoint if/ when the two modes diverge.\n",
        "- Create `tf.Variable`s\n",
        "outside the Python function and modify them on the inside. The same goes for objects that use `tf.Variable`, like `tf.keras.layers`, `tf.keras.Model`s and `tf.keras.optimizers`.\n",
        "- Avoid writing functions that depend on outer Python variables, excluding `tf.Variable`s and Keras objects. Learn more in _Depending on Python global and free variables_ of the [`tf.function` guide](./function.ipynb).\n",
        "- Prefer to write functions which take tensors and other TensorFlow types as input. You can pass in other object types but be careful! Learn more in _Depending on Python objects_ of the [`tf.function` guide](./function.ipynb).\n",
        "- Include as much computation as possible under a `tf.function` to maximize the performance gain. For example, decorate a whole training step or the entire training loop.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ViM3oBJVJrDx"
      },
      "source": [
        "## Seeing the speed-up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A6NHDp7vAKcJ"
      },
      "source": [
        "`tf.function` usually improves the performance of your code, but the amount of speed-up depends on the kind of computation you run. Small computations can be dominated by the overhead of calling a graph. You can measure the difference in performance like so:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jr7p1BBjauPK"
      },
      "outputs": [],
      "source": [
        "x = tf.random.uniform(shape=[10, 10], minval=-1, maxval=2, dtype=tf.dtypes.int32)\n",
        "\n",
        "def power(x, y):\n",
        "  result = tf.eye(10, dtype=tf.dtypes.int32)\n",
        "  for _ in range(y):\n",
        "    result = tf.matmul(x, result)\n",
        "  return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ms2yJyAnUYxK"
      },
      "outputs": [],
      "source": [
        "print(\"Eager execution:\", timeit.timeit(lambda: power(x, 100), number=1000), \"seconds\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gUB2mTyRYRAe"
      },
      "outputs": [],
      "source": [
        "power_as_graph = tf.function(power)\n",
        "print(\"Graph execution:\", timeit.timeit(lambda: power_as_graph(x, 100), number=1000), \"seconds\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q1Pfo5YwwILi"
      },
      "source": [
        "`tf.function` is commonly used to speed up training loops, and you can learn more about it in the _Speeding-up your training step with `tf.function`_ section of the [Writing a training loop from scratch](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch) with Keras guide.\n",
        "\n",
        "Note: You can also try `tf.function(jit_compile=True)` for a more significant performance boost, especially if your code is heavy on TensorFlow control flow and uses many small tensors. Learn more in the _Explicit compilation with `tf.function(jit_compile=True)`_ section of the [XLA overview](https://www.tensorflow.org/xla)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sm0bNFp8PX53"
      },
      "source": [
        "### Performance and trade-offs\n",
        "\n",
        "Graphs can speed up your code, but the process of creating them has some overhead. For some functions, the creation of the graph takes more time than the execution of the graph. **This investment is usually quickly paid back with the performance boost of subsequent executions, but it's important to be aware that the first few  steps of any large model training can be slower due to tracing.**\n",
        "\n",
        "No matter how large your model, you want to avoid tracing frequently. In the _Controlling retracing_ section, the [`tf.function` guide](./function.ipynb) discusses how to set input specifications and use tensor arguments to avoid retracing. If you find you are getting unusually poor performance, it's a good idea to check if you are retracing accidentally."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F4InDaTjwmBA"
      },
      "source": [
        "## When is a `tf.function` tracing?\n",
        "\n",
        "To figure out when your `tf.function` is tracing, add a `print` statement to its code. As a rule of thumb, `tf.function` will execute the `print` statement every time it traces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hXtwlbpofLgW"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def a_function_with_python_side_effect(x):\n",
        "  print(\"Tracing!\") # An eager-only side effect.\n",
        "  return x * x + tf.constant(2)\n",
        "\n",
        "# This is traced the first time.\n",
        "print(a_function_with_python_side_effect(tf.constant(2)))\n",
        "# The second time through, you won't see the side effect.\n",
        "print(a_function_with_python_side_effect(tf.constant(3)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "inzSg8yzfNjl"
      },
      "outputs": [],
      "source": [
        "# This retraces each time the Python argument changes,\n",
        "# as a Python argument could be an epoch count or other\n",
        "# hyperparameter.\n",
        "print(a_function_with_python_side_effect(2))\n",
        "print(a_function_with_python_side_effect(3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rtN8NW6AfKye"
      },
      "source": [
        "New Python arguments always trigger the creation of a new graph, hence the extra tracing.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D1kbr5ocpS6R"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "You can learn more about `tf.function` on the API reference page and by following the [Better performance with `tf.function`](./function.ipynb) guide."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_to_graphs.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}