add comment

notebooks/thunder_trace_intro.ipynb

@@ -411,15 +411,6 @@
    "cell_type": "code",
    "execution_count": 14,
    "metadata": {},
-   "outputs": [],
-   "source": [
-    "trc = Trace()  # Create an empty trace"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
    "outputs": [
     {
      "data": {
@@ -432,15 +423,17 @@
        "# No signature available"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
+    "tmp_trc = Trace()  # Create an empty trace\n",
+    "\n",
     "# Print the trace.\n",
     "# There is nothing interesting in it.\n",
-    "trc"
+    "tmp_trc"
    ]
   },
   {
@@ -449,7 +442,7 @@
    "source": [
     "##### Adding a signature for the trace\n",
     "\n",
-    "As printed above, we can see that our trace is very empty. First, we will add signature for the function that trace will represent.\n",
+    "As printed above, we can see that our trace is empty and doesn't have a signature for the function it represents. First, we will add signature for the function that trace will represent.\n",
     "\n",
     "There are two ways that we can specify what the signature of the trace will be.\n",
     "\n",
@@ -459,7 +452,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -473,7 +466,7 @@
        "def add_fn(x, y):"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -495,7 +488,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
@@ -509,7 +502,7 @@
        "def my_trace(x, y):"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -540,10 +533,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Create a new Trace object.\n",
+    "trc = Trace()\n",
+    "\n",
     "# Set tracing context to `trc` trace.\n",
     "# Any information related to trace, will be looked from or added to `trc` trace.\n",
     "with tracectx(trc):\n",
@@ -564,7 +560,7 @@
     "    return_bsym = python_return.bind(o, output=None)\n",
     "\n",
     "\n",
-    "# Set the sig-info\n",
+    "# Set the signature for Trace\n",
     "trc_args = (x, y)\n",
     "\n",
     "si = SigInfo(\"my_trace\")\n",
@@ -579,7 +575,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
@@ -592,12 +588,12 @@
        "\n",
        "@torch.no_grad()\n",
        "@no_autocast\n",
-       "def my_trace(t2, t3):\n",
-       "  t5 = ltorch.add(t2, t3, alpha=1)  # t5: \"cpu f32[3, 3, 3]\"\n",
-       "  return t5"
+       "def my_trace(x, y):\n",
+       "  t1 = ltorch.add(x, y, alpha=1)  # t1: \"cpu f32[3, 3, 3]\"\n",
+       "  return t1"
       ]
      },
-     "execution_count": 21,
+     "execution_count": 18,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -634,6 +630,9 @@
     "    # internally calls `bind` on the symbol with args, kwargs and output\n",
     "    # and the generated BoundSymbol is added to the current trace.\n",
     "    # We did this manually in the above approach.\n",
+    "    # NOTE: `ltorch.add` is composed of `prims.broadcast_in_dims` (for broadcasting inputs)\n",
+    "    #       and `prims.add` symbols. So, when we call into `ltorch.add`, it also traces the mentioned symbols and\n",
+    "    #       they are added in the trace indented below `ltorch.add`.\n",
     "    o = ltorch.add(x, y)\n",
     "    python_return(o)\n",
     "\n",
@@ -646,7 +645,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -666,7 +665,7 @@
        "  return t1"
       ]
      },
-     "execution_count": 19,
+     "execution_count": 20,
      "metadata": {},
      "output_type": "execute_result"
     }