TF Autograph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hrXv0rU9sIma"
   },
   "source": [
    "# Autograph: Basic\n",
    "In this ungraded lab, you will go through  some of the basics of autograph so you can explore what the generated code looks like."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NiolgWMPgpwI"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Addition in autograph\n",
    "You can use the `@tf.function` decorator to automatically generate the graph-style code as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[5. 2.]\n",
      " [4. 9.]], shape=(2, 2), dtype=float32)\n",
      "def tf__add(a, b):\n",
      "    with ag__.FunctionScope('add', '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",
      "        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": [
    "@tf.function\n",
    "def add(a, b):\n",
    "    return a + b\n",
    "\n",
    "\n",
    "a = tf.Variable([[1.,2.],[3.,4.]])\n",
    "b = tf.Variable([[4.,0.],[1.,5.]])\n",
    "print(tf.add(a, b))\n",
    "\n",
    "# See what the generated code looks like\n",
    "print(tf.autograph.to_code(add.python_function))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 921 µs, sys: 0 ns, total: 921 µs\n",
      "Wall time: 643 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2), dtype=float32, numpy=\n",
       "array([[5., 2.],\n",
       "       [4., 9.]], dtype=float32)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "add(a,b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[5. 2.]\n",
      " [4. 9.]], shape=(2, 2), dtype=float32)\n",
      "CPU times: user 4.12 ms, sys: 34 µs, total: 4.16 ms\n",
      "Wall time: 4.13 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "def add(a, b):\n",
    "    return a + b\n",
    "\n",
    "a = tf.Variable([[1.,2.],[3.,4.]])\n",
    "b = tf.Variable([[4.,0.],[1.,5.]])\n",
    "print(add(a, b))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## if-statements in autograph\n",
    "Control flow statements which are very intuitive to write in eager mode can look very complex in graph mode. You can see that in the next examples: first a simple function, then a more complicated one that involves lots of ops and conditionals (fizzbuzz)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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 if_body():\n",
      "            nonlocal x\n",
      "            x = (ag__.ld(x) * ag__.ld(x))\n",
      "\n",
      "        def else_body():\n",
      "            nonlocal x\n",
      "            pass\n",
      "        ag__.if_stmt((ag__.ld(x) > 0), if_body, else_body, get_state, set_state, ('x',), 1)\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": [
    "# simple function that returns the square if the input is greater than zero\n",
    "@tf.function\n",
    "def f(x):\n",
    "    if x>0:\n",
    "        x = x * x\n",
    "    return x\n",
    "\n",
    "print(tf.autograph.to_code(f.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fizzbuzz in autograph\n",
    "\n",
    "You may remember implementing [fizzbuzz](http://wiki.c2.com/?FizzBuzzTest) in preparation for a coding interview.  \n",
    "- Imagine how much fun it would be if you were asked to impement the graph mode version of that code!  \n",
    "\n",
    "Fortunately, you can just use `@tf.function` and then call `tf.autograph.to_code`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def tf__fizzbuzz(max_num):\n",
      "    with ag__.FunctionScope('fizzbuzz', '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",
      "        counter = 0\n",
      "\n",
      "        def get_state_3():\n",
      "            return (counter,)\n",
      "\n",
      "        def set_state_3(vars_):\n",
      "            nonlocal counter\n",
      "            (counter,) = vars_\n",
      "\n",
      "        def loop_body(itr):\n",
      "            nonlocal counter\n",
      "            num = itr\n",
      "\n",
      "            def get_state_2():\n",
      "                return ()\n",
      "\n",
      "            def set_state_2(block_vars):\n",
      "                pass\n",
      "\n",
      "            def if_body_2():\n",
      "                ag__.ld(print)('FizzBuzz')\n",
      "\n",
      "            def else_body_2():\n",
      "\n",
      "                def get_state_1():\n",
      "                    return ()\n",
      "\n",
      "                def set_state_1(block_vars):\n",
      "                    pass\n",
      "\n",
      "                def if_body_1():\n",
      "                    ag__.ld(print)('Fizz')\n",
      "\n",
      "                def else_body_1():\n",
      "\n",
      "                    def get_state():\n",
      "                        return ()\n",
      "\n",
      "                    def set_state(block_vars):\n",
      "                        pass\n",
      "\n",
      "                    def if_body():\n",
      "                        ag__.ld(print)('Buzz')\n",
      "\n",
      "                    def else_body():\n",
      "                        ag__.ld(print)(ag__.ld(num))\n",
      "                    ag__.if_stmt(((ag__.ld(num) % 5) == 0), if_body, else_body, get_state, set_state, (), 0)\n",
      "                ag__.if_stmt(((ag__.ld(num) % 3) == 0), if_body_1, else_body_1, get_state_1, set_state_1, (), 0)\n",
      "            ag__.if_stmt(ag__.and_((lambda : ((ag__.ld(num) % 3) == 0)), (lambda : ((ag__.ld(num) % 5) == 0))), if_body_2, else_body_2, get_state_2, set_state_2, (), 0)\n",
      "            counter = ag__.ld(counter)\n",
      "            counter += 1\n",
      "        num = ag__.Undefined('num')\n",
      "        ag__.for_stmt(ag__.converted_call(ag__.ld(range), (ag__.ld(max_num),), None, fscope), None, loop_body, get_state_3, set_state_3, ('counter',), {'iterate_names': 'num'})\n",
      "        try:\n",
      "            do_return = True\n",
      "            retval_ = ag__.ld(counter)\n",
      "        except:\n",
      "            do_return = False\n",
      "            raise\n",
      "        return fscope.ret(retval_, do_return)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def fizzbuzz(max_num):\n",
    "    counter = 0\n",
    "    for num in range(max_num):\n",
    "        if num % 3 == 0 and num % 5 == 0:\n",
    "            print('FizzBuzz')\n",
    "        elif num % 3 == 0:\n",
    "            print('Fizz')\n",
    "        elif num % 5 == 0:\n",
    "            print('Buzz')\n",
    "        else:\n",
    "            print(num)\n",
    "        counter += 1\n",
    "    return counter\n",
    "\n",
    "print(tf.autograph.to_code(fizzbuzz.python_function))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Autograph Basics.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
}