Automated rollback of commit 49eca1f

PiperOrigin-RevId: 405918182

Author: tensorflower-gardener

site/en/guide/function.ipynb

@@ -927,7 +927,7 @@
         "    state = rnn_step(input_data[i], state)\n",
         "    states = states.write(i, state)\n",
         "  return tf.transpose(states.stack(), [1, 0, 2])\n",
-        "  \n",
+        "\n",
         "dynamic_rnn(rnn_step,\n",
         "            tf.random.uniform([batch_size, seq_len, feature_size]),\n",
         "            tf.zeros([batch_size, feature_size]))"
@@ -1017,13 +1017,86 @@
         "assert len(external_list) == 1"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5eZTFRv_k_nR"
+      },
+      "source": [
+        "Sometimes unexpected behaviors are very hard to notice. In the example below, the `counter` is intended to safeguard the increment of a variable. However because it is a python integer and not a TensorFlow object, it's value is captured during the first trace. When the `tf.function` is used, the `assign_add` will be recorded unconditionally in the underlying graph. Therefore `v` will increase by 1, every time the `tf.function` is called. This issue is common among users that try to migrate their Grpah-mode Tensorflow code to Tensorflow 2 using `tf.function` decorators, when python side-effects (the `counter` in the example) are used to determine what ops to run (`assign_add` in the example). Usually, users realize this only after seeing suspicious numerical results, or significantly lower performance than expected (e.g. if the guarded operation is very costly)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5r6p7-9jk_3L"
+      },
+      "outputs": [],
+      "source": [
+        "class Model(tf.Module):\n",
+        "  def __init__(self):\n",
+        "    self.v = tf.Variable(0)\n",
+        "    self.counter = 0\n",
+        "\n",
+        "  @tf.function\n",
+        "  def __call__(self):\n",
+        "    if self.counter == 0:\n",
+        "      # A python side-effect\n",
+        "      self.counter += 1\n",
+        "      self.v.assign_add(1)\n",
+        "\n",
+        "    return self.v\n",
+        "\n",
+        "m = Model()\n",
+        "for n in range(3):\n",
+        "  print(m().numpy()) # prints 1, 2, 3"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tXCTcHoVcxhX"
+      },
+      "source": [
+        "A workaround to achieve the expected behavior is using [`tf.init_scope`](https://www.tensorflow.org/api_docs/python/tf/init_scope) to lift the operations outside of the function graph. This ensures that the variable increment is only done once during tracing time. It should be noted `init_scope` has other side effects including cleared control flow and gradient tape. Sometimes the usage of `init_scope` can become too complex to manage realistically."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "An4MrIbrcvi8"
+      },
+      "outputs": [],
+      "source": [
+        "class Model(tf.Module):\n",
+        "  def __init__(self):\n",
+        "    self.v = tf.Variable(0)\n",
+        "    self.counter = 0\n",
+        "\n",
+        "  @tf.function\n",
+        "  def __call__(self):\n",
+        "    if self.counter == 0:\n",
+        "      # Lifts ops out of function-building graphs\n",
+        "      with tf.init_scope():\n",
+        "        self.counter += 1\n",
+        "        self.v.assign_add(1)\n",
+        "\n",
+        "    return self.v\n",
+        "\n",
+        "m = Model()\n",
+        "for n in range(3):\n",
+        "  print(m().numpy()) # prints 1, 1, 1"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "pbFG5CX4LwQA"
       },
       "source": [
-        "You should avoid mutating containers like lists, dicts, other objects that live outside the `Function`. Instead, use arguments and TF objects. For example, the section [\"Accumulating values in a loop\"](#accumulating_values_in_a_loop) has one example of how list-like operations can be implemented.\n",
+        "In summary, as a rule of thumb, you should avoid mutating python objects such as integers or containers like lists that live outside the `Function`. Instead, use arguments and TF objects. For example, the section [\"Accumulating values in a loop\"](#accumulating_values_in_a_loop) has one example of how list-like operations can be implemented.\n",
         "\n",
         "You can, in some cases, capture and manipulate state if it is a [`tf.Variable`](https://www.tensorflow.org/guide/variable). This is how the weights of Keras models are updated with repeated calls to the same `ConcreteFunction`."
       ]
@@ -1625,6 +1698,7 @@
     "colab": {
       "collapsed_sections": [],
       "name": "function.ipynb",
+      "provenance": [],
       "toc_visible": true
     },
     "kernelspec": {