diff --git a/docs/website/docs/extensions/index.md b/docs/website/docs/extensions/index.md
index 6992cc3..811cf37 100644
--- a/docs/website/docs/extensions/index.md
+++ b/docs/website/docs/extensions/index.md
@@ -167,14 +167,14 @@
 ```mlir
 // Main user module compiled by IREE:
 module @model {
-  func @predict() {
+  func.func @predict() {
     %4 = call @my_custom_module::@some_func(%3) : (tensor<?xf32>) -> i32
     ...
   }
 }
 // External module that will be available at runtime:
 module @my_custom_module {
-  func @some_func(%input: tensor<?xf32>) -> i32  // note empty for extern
+  func.func @some_func(%input: tensor<?xf32>) -> i32  // note empty for extern
 }
 ```
 
diff --git a/samples/colab/low_level_invoke_function.ipynb b/samples/colab/low_level_invoke_function.ipynb
index e16a2ca..62fafec 100644
--- a/samples/colab/low_level_invoke_function.ipynb
+++ b/samples/colab/low_level_invoke_function.ipynb
@@ -63,7 +63,7 @@
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "outputId": "bfd1cc6a-bc75-4d16-850d-64776dc6d64f"
+        "outputId": "b3810038-9b88-48b8-d3dc-1d8c37300ac8"
       },
       "source": [
         "!python -m pip install iree-compiler iree-runtime -f https://github.com/google/iree/releases"
@@ -72,18 +72,20 @@
       "outputs": [
         {
           "output_type": "stream",
+          "name": "stdout",
           "text": [
             "Looking in links: https://github.com/google/iree/releases\n",
             "Collecting iree-compiler\n",
-            "\u001b[?25l  Downloading https://github.com/google/iree/releases/download/snapshot-20210608.328/iree_compiler_snapshot-20210608.328-py3-none-manylinux2014_x86_64.whl (33.6MB)\n",
-            "\u001b[K     |████████████████████████████████| 33.6MB 118kB/s \n",
+            "  Downloading https://github.com/google/iree/releases/download/candidate-20220514.138/iree_compiler-20220514.138-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (49.0 MB)\n",
+            "\u001b[K     |████████████████████████████████| 49.0 MB 8.9 kB/s \n",
             "\u001b[?25hCollecting iree-runtime\n",
-            "\u001b[?25l  Downloading https://github.com/google/iree/releases/download/snapshot-20210608.328/iree_runtime_snapshot-20210608.328-cp37-cp37m-manylinux2014_x86_64.whl (521kB)\n",
-            "\u001b[K     |████████████████████████████████| 522kB 27.1MB/s \n",
-            "\u001b[?25hInstalling collected packages: iree-compiler, iree-runtime\n",
-            "Successfully installed iree-compiler-20210608.328 iree-runtime-20210608.328\n"
-          ],
-          "name": "stdout"
+            "  Downloading https://github.com/google/iree/releases/download/candidate-20220514.138/iree_runtime-20220514.138-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.2 MB)\n",
+            "\u001b[K     |████████████████████████████████| 2.2 MB 40.2 MB/s \n",
+            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from iree-compiler) (1.21.6)\n",
+            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from iree-compiler) (3.13)\n",
+            "Installing collected packages: iree-runtime, iree-compiler\n",
+            "Successfully installed iree-compiler-20220514.138 iree-runtime-20220514.138\n"
+          ]
         }
       ]
     },
@@ -110,10 +112,9 @@
         "# Compile a module.\n",
         "SIMPLE_MUL_ASM = \"\"\"\n",
         "  module @arithmetic {\n",
-        "    func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>\n",
-        "          {\n",
-        "        %0 = arith.mulf %arg0, %arg1 : tensor<4xf32>\n",
-        "        return %0 : tensor<4xf32>\n",
+        "    func.func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {\n",
+        "      %0 = arith.mulf %arg0, %arg1 : tensor<4xf32>\n",
+        "      return %0 : tensor<4xf32>\n",
         "    } \n",
         "  }\n",
         "\"\"\"\n",
@@ -121,7 +122,7 @@
         "compiled_flatbuffer = compile_str(SIMPLE_MUL_ASM, target_backends=[\"vmvx\"])\n",
         "vm_module = ireert.VmModule.from_flatbuffer(compiled_flatbuffer)"
       ],
-      "execution_count": 7,
+      "execution_count": 4,
       "outputs": []
     },
     {
@@ -131,7 +132,7 @@
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "outputId": "37e2df44-90fb-425e-f6e5-059efdfe099e"
+        "outputId": "933b2530-b3f9-4dbf-bdfd-8d222843e2cf"
       },
       "source": [
         "# Register the module with a runtime context.\n",
@@ -145,28 +146,20 @@
         "arg0 = np.array([1., 2., 3., 4.], dtype=np.float32)\n",
         "arg1 = np.array([4., 5., 6., 7.], dtype=np.float32)\n",
         "f = ctx.modules.arithmetic[\"simple_mul\"]\n",
-        "results = f(arg0, arg1)\n",
+        "results = f(arg0, arg1).to_host()\n",
         "print(\"Results:\", results)"
       ],
-      "execution_count": 8,
+      "execution_count": 6,
       "outputs": [
         {
           "output_type": "stream",
+          "name": "stdout",
           "text": [
             "INVOKE simple_mul\n",
             "Results: [ 4. 10. 18. 28.]\n"
-          ],
-          "name": "stdout"
-        },
-        {
-          "output_type": "stream",
-          "text": [
-            "Created IREE driver vmvx: <iree.runtime.binding.HalDriver object at 0x7f9b9d6e6ef0>\n",
-            "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7f9b9d6e6ef0>\n"
-          ],
-          "name": "stderr"
+          ]
         }
       ]
     }
   ]
-}
+}
\ No newline at end of file