Fix function.ipynb

PiperOrigin-RevId: 547932510

Author: MarkDaoust

site/en/guide/function.ipynb

@@ -91,8 +91,6 @@
       },
       "outputs": [],
       "source": [
-        "# Update TensorFlow, as this notebook requires version 2.9 or later\n",
-        "!pip install -q -U tensorflow>=2.9.0\n",
         "import tensorflow as tf"
       ]
     },
@@ -362,7 +360,7 @@
       "source": [
         "#### Rules of tracing\n",
         "\n",
-        "When called, a `Function` matches the call arguments to existing `ConcreteFunction`s using `tf.types.experimental.TraceType` of each argument. If a matching `ConcreteFunction` is found, the call is dispatched to it. If no match is found, a new `ConcreteFunction` is traced. \n",
+        "When called, a `Function` matches the call arguments to existing `ConcreteFunction`s using `tf.types.experimental.TraceType` of each argument. If a matching `ConcreteFunction` is found, the call is dispatched to it. If no match is found, a new `ConcreteFunction` is traced.\n",
         "\n",
         "If multiple matches are found, the most specific signature is chosen. Matching is done by [subtyping](https://en.wikipedia.org/wiki/Subtyping), much like normal function calls in C++ or Java, for instance. For example, `TensorShape([1, 2])` is a subtype of `TensorShape([None, None])` and so a call to the tf.function with `TensorShape([1, 2])` can be dispatched to the `ConcreteFunction` produced with `TensorShape([None, None])` but if a `ConcreteFunction` with `TensorShape([1, None])` also exists then it will prioritized since it is more specific.\n",
         "\n",
@@ -422,11 +420,11 @@
         "\n",
         "print(next_collatz(tf.constant([1, 2])))\n",
         "# You specified a 1-D tensor in the input signature, so this should fail.\n",
-        "with assert_raises(ValueError):\n",
+        "with assert_raises(TypeError):\n",
         "  next_collatz(tf.constant([[1, 2], [3, 4]]))\n",
         "\n",
         "# You specified an int32 dtype in the input signature, so this should fail.\n",
-        "with assert_raises(ValueError):\n",
+        "with assert_raises(TypeError):\n",
         "  next_collatz(tf.constant([1.0, 2.0]))\n"
       ]
     },
@@ -560,8 +558,8 @@
         "  flavor = tf.constant([3, 4])\n",
         "\n",
         "# As described in the above rules, a generic TraceType for `Apple` and `Mango`\n",
-        "# is generated (and a corresponding ConcreteFunction is traced) but it fails to \n",
-        "# match the second function call since the first pair of Apple() and Mango() \n",
+        "# is generated (and a corresponding ConcreteFunction is traced) but it fails to\n",
+        "# match the second function call since the first pair of Apple() and Mango()\n",
         "# have gone out out of scope by then and deleted.\n",
         "get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function\n",
         "get_mixed_flavor(Apple(), Mango()) # Traces a new concrete function again\n",
@@ -591,7 +589,7 @@
         "\n",
         "  def __eq__(self, other):\n",
         "    return type(other) is FruitTraceType and self.fruit_type == other.fruit_type\n",
-        "  \n",
+        "\n",
         "  def __hash__(self):\n",
         "    return hash(self.fruit_type)\n",
         "\n",
@@ -970,7 +968,7 @@
         "id": "JeD2U-yrbfVb"
       },
       "source": [
-        "When wrapping Python/NumPy data in a Dataset, be mindful of `tf.data.Dataset.from_generator` versus ` tf.data.Dataset.from_tensors`. The former will keep the data in Python and fetch it via `tf.py_function` which can have performance implications, whereas the latter will bundle a copy of the data as one large `tf.constant()` node in the graph, which can have memory implications.\n",
+        "When wrapping Python/NumPy data in a Dataset, be mindful of `tf.data.Dataset.from_generator` versus ` tf.data.Dataset.from_tensor_slices`. The former will keep the data in Python and fetch it via `tf.py_function` which can have performance implications, whereas the latter will bundle a copy of the data as one large `tf.constant()` node in the graph, which can have memory implications.\n",
         "\n",
         "Reading data from files via `TFRecordDataset`, `CsvDataset`, etc. is the most effective way to consume data, as then TensorFlow itself can manage the asynchronous loading and prefetching of data, without having to involve Python. To learn more, see the [`tf.data`: Build TensorFlow input pipelines](../../guide/data) guide."
       ]
@@ -1608,7 +1606,7 @@
         "new_model = SimpleModel()\n",
         "evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)\n",
         "# Don't pass in `new_model`, `Function` already captured its state during tracing.\n",
-        "print(evaluate_no_bias(x))  "
+        "print(evaluate_no_bias(x))"
       ]
     },
     {
@@ -1752,7 +1750,7 @@
       "source": [
         "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
         "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
-        " \n",
+        "\n",
         "@tf.function\n",
         "def train_step(w, x, y, optimizer):\n",
         "   with tf.GradientTape() as tape:\n",
@@ -1802,13 +1800,13 @@
         "y = tf.constant([2.])\n",
         "\n",
         "# Make a new Function and ConcreteFunction for each optimizer.\n",
-        "train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)\n",
-        "train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)\n",
+        "train_step_1 = tf.function(train_step)\n",
+        "train_step_2 = tf.function(train_step)\n",
         "for i in range(10):\n",
         "  if i % 2 == 0:\n",
-        "    train_step_1(w, x, y) # `opt1` is not used as a parameter. \n",
+        "    train_step_1(w, x, y, opt1)\n",
         "  else:\n",
-        "    train_step_2(w, x, y) # `opt2` is not used as a parameter."
+        "    train_step_2(w, x, y, opt2)"
       ]
     },
     {