C2_W3_Lab_2-graphs-for-complex-code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/https-deeplearning-ai/tensorflow-3-public/blob/main/Course%202%20-%20Custom%20Training%20loops%2C%20Gradients%20and%20Distributed%20Training/Week%203%20-%20Autograph/C2_W3_Lab_2-graphs-for-complex-code.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hrXv0rU9sIma"
   },
   "source": [
    "# Autograph: Graphs for complex code\n",
    "\n",
    "In this ungraded lab, you'll go through some of the scenarios from the lesson `Creating graphs for complex code`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NiolgWMPgpwI"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    # %tensorflow_version only exists in Colab.\n",
    "    %tensorflow_version 2.x\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you saw in the lectures, seemingly simple functions can sometimes be difficult to write in graph mode. Fortunately, Autograph generates this complex graph code for us.\n",
    "\n",
    "- Here is a function that does some multiplication and additon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(10.0, shape=(), dtype=float32)\n",
      "def tf__f(x, y):\n",
      "    with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "        do_return = False\n",
      "        retval_ = ag__.UndefinedReturnValue()\n",
      "        ag__.converted_call(ag__.ld(a).assign, ((ag__.ld(y) * ag__.ld(b)),), None, fscope)\n",
      "        ag__.converted_call(ag__.ld(b).assign_add, ((ag__.ld(x) * ag__.ld(a)),), None, fscope)\n",
      "        try:\n",
      "            do_return = True\n",
      "            retval_ = (ag__.ld(a) + ag__.ld(b))\n",
      "        except:\n",
      "            do_return = False\n",
      "            raise\n",
      "        return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "a = tf.Variable(1.0)\n",
    "b = tf.Variable(2.0)\n",
    "\n",
    "@tf.function\n",
    "def f(x,y):\n",
    "    a.assign(y * b)\n",
    "    b.assign_add(x * a)\n",
    "    return a + b\n",
    "\n",
    "print(f(1.0, 2.0))\n",
    "\n",
    "print(tf.autograph.to_code(f.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is a function that checks if the sign of a number is positive or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sign = b'Positive'\n",
      "Sign = b'Negative or zero'\n",
      "def tf__sign(x):\n",
      "    with ag__.FunctionScope('sign', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "        do_return = False\n",
      "        retval_ = ag__.UndefinedReturnValue()\n",
      "\n",
      "        def get_state():\n",
      "            return (do_return, retval_)\n",
      "\n",
      "        def set_state(vars_):\n",
      "            nonlocal do_return, retval_\n",
      "            (do_return, retval_) = vars_\n",
      "\n",
      "        def if_body():\n",
      "            nonlocal do_return, retval_\n",
      "            try:\n",
      "                do_return = True\n",
      "                retval_ = 'Positive'\n",
      "            except:\n",
      "                do_return = False\n",
      "                raise\n",
      "\n",
      "        def else_body():\n",
      "            nonlocal do_return, retval_\n",
      "            try:\n",
      "                do_return = True\n",
      "                retval_ = 'Negative or zero'\n",
      "            except:\n",
      "                do_return = False\n",
      "                raise\n",
      "        ag__.if_stmt((ag__.ld(x) > 0), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)\n",
      "        return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def sign(x):\n",
    "    if x > 0:\n",
    "        return 'Positive'\n",
    "    else:\n",
    "        return 'Negative or zero'\n",
    "\n",
    "print(\"Sign = {}\".format(sign(tf.constant(2))))\n",
    "print(\"Sign = {}\".format(sign(tf.constant(-2))))\n",
    "\n",
    "print(tf.autograph.to_code(sign.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is another function that includes a while loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def tf__f(x):\n",
      "    with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "        do_return = False\n",
      "        retval_ = ag__.UndefinedReturnValue()\n",
      "\n",
      "        def get_state():\n",
      "            return (x,)\n",
      "\n",
      "        def set_state(vars_):\n",
      "            nonlocal x\n",
      "            (x,) = vars_\n",
      "\n",
      "        def loop_body():\n",
      "            nonlocal x\n",
      "            ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope)\n",
      "            x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope)\n",
      "\n",
      "        def loop_test():\n",
      "            return (ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1)\n",
      "        ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {})\n",
      "        try:\n",
      "            do_return = True\n",
      "            retval_ = ag__.ld(x)\n",
      "        except:\n",
      "            do_return = False\n",
      "            raise\n",
      "        return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    while tf.reduce_sum(x) > 1:\n",
    "        tf.print(x)\n",
    "        x = tf.tanh(x)\n",
    "    return x\n",
    "\n",
    "print(tf.autograph.to_code(f.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Here is a function that uses a for loop and an if statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def tf__sum_even(items):\n",
      "    with ag__.FunctionScope('sum_even', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "        do_return = False\n",
      "        retval_ = ag__.UndefinedReturnValue()\n",
      "        s = 0\n",
      "\n",
      "        def get_state_2():\n",
      "            return (s,)\n",
      "\n",
      "        def set_state_2(vars_):\n",
      "            nonlocal s\n",
      "            (s,) = vars_\n",
      "\n",
      "        def loop_body(itr):\n",
      "            nonlocal s\n",
      "            c = itr\n",
      "            continue_ = False\n",
      "\n",
      "            def get_state():\n",
      "                return (continue_,)\n",
      "\n",
      "            def set_state(vars_):\n",
      "                nonlocal continue_\n",
      "                (continue_,) = vars_\n",
      "\n",
      "            def if_body():\n",
      "                nonlocal continue_\n",
      "                continue_ = True\n",
      "\n",
      "            def else_body():\n",
      "                nonlocal continue_\n",
      "                pass\n",
      "            ag__.if_stmt(((ag__.ld(c) % 2) > 0), if_body, else_body, get_state, set_state, ('continue_',), 1)\n",
      "\n",
      "            def get_state_1():\n",
      "                return (s,)\n",
      "\n",
      "            def set_state_1(vars_):\n",
      "                nonlocal s\n",
      "                (s,) = vars_\n",
      "\n",
      "            def if_body_1():\n",
      "                nonlocal s\n",
      "                s = ag__.ld(s)\n",
      "                s += c\n",
      "\n",
      "            def else_body_1():\n",
      "                nonlocal s\n",
      "                pass\n",
      "            ag__.if_stmt(ag__.not_(continue_), if_body_1, else_body_1, get_state_1, set_state_1, ('s',), 1)\n",
      "        c = ag__.Undefined('c')\n",
      "        continue_ = ag__.Undefined('continue_')\n",
      "        ag__.for_stmt(ag__.ld(items), None, loop_body, get_state_2, set_state_2, ('s',), {'iterate_names': 'c'})\n",
      "        try:\n",
      "            do_return = True\n",
      "            retval_ = ag__.ld(s)\n",
      "        except:\n",
      "            do_return = False\n",
      "            raise\n",
      "        return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def sum_even(items):\n",
    "    s = 0\n",
    "    for c in items:\n",
    "        if c % 2 > 0:\n",
    "            continue\n",
    "        s += c\n",
    "    return s\n",
    "\n",
    "print(tf.autograph.to_code(sum_even.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Print statements\n",
    "\n",
    "Tracing also behaves differently in graph mode. First, here is a function (not decorated with `@tf.function` yet) that prints the value of the input parameter.  `f(2)` is called in a for loop 5 times, and then `f(3)` is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traced with 2\n",
      "Traced with 2\n",
      "Traced with 2\n",
      "Traced with 2\n",
      "Traced with 2\n",
      "Traced with 3\n"
     ]
    }
   ],
   "source": [
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you were to decorate this function with `@tf.function` and run it, notice that the print statement only appears once for `f(2)` even though it is called in a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traced with 2\n",
      "Traced with 3\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now compare `print` to `tf.print`.\n",
    "- `tf.print` is graph aware and will run as expected in loops. \n",
    "\n",
    "Try running the same code where `tf.print()` is added in addition to the regular `print`.\n",
    "- Note how `tf.print` behaves compared to `print` in graph mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traced with 2\n",
      "Executed with 2\n",
      "Executed with 2\n",
      "Executed with 2\n",
      "Executed with 2\n",
      "Executed with 2\n",
      "Traced with 3\n",
      "Executed with 3\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    print(\"Traced with\", x)\n",
    "    # added tf.print\n",
    "    tf.print(\"Executed with\", x)\n",
    "\n",
    "for i in range(5):\n",
    "    f(2)\n",
    "    \n",
    "f(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Avoid defining variables inside the function\n",
    "\n",
    "This function (not decorated yet) defines a tensor `v` and adds the input `x` to it.  \n",
    "\n",
    "Here, it runs fine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>\n"
     ]
    }
   ],
   "source": [
    "def f(x):\n",
    "    v = tf.Variable(1.0)\n",
    "    v.assign_add(x)\n",
    "    return v\n",
    "\n",
    "print(f(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if you decorate the function with `@tf.function`.\n",
    "\n",
    "The cell below will throw an error because `tf.Variable` is defined within the function. The graph mode function should only contain operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(2.0, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "v = tf.Variable(1.0)\n",
    "@tf.function\n",
    "def f(x):\n",
    "    \n",
    "    v.assign_add(x)\n",
    "    return v\n",
    "\n",
    "print(f(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get around the error above, simply move `v = tf.Variable(1.0)` to the top of the cell before the `@tf.function` decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(6.0, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# define the variables outside of the decorated function\n",
    "v = tf.Variable(1.0)\n",
    "\n",
    "@tf.function\n",
    "def f(x):\n",
    "    return v.assign_add(x)\n",
    "\n",
    "print(f(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Autograph Lesson 2.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}