Externalization Documentaion Updates

coreai_torch/converter.py

@@ -310,6 +310,14 @@ def add_pytorch_module(
         ``to_coreai``).
 
         Args:
+            model: The PyTorch module to convert.
+            export_fn: Callable that takes ``model`` and returns its decomposed
+                ``ExportedProgram`` (see above). It is called once when staging,
+                and again during ``to_coreai`` if ``externalize_modules`` is set.
+            externalize_modules: Submodule classes to keep as separate graphs,
+                given as :class:`ExternalizeSpec` objects (composite ops) or bare
+                classes (experimental simple externalization). Only submodules
+                of ``model`` are matched, never ``model`` itself.
             input_names: Non-stateful forward() arg names. See
                 :meth:`add_exported_program`.
             output_names: Return value names. See :meth:`add_exported_program`.
@@ -318,8 +326,16 @@ def add_pytorch_module(
                 :meth:`add_exported_program`.
             output_image_types: Outputs to treat as images. See
                 :meth:`add_exported_program`.
+            entrypoint_name: Name of this program's entry-point graph in the
+                converted module. Must be unique among staged programs.
 
-        Returns ``self`` for chaining.
+        Returns:
+            This converter, so further calls can be chained onto it, e.g.
+            ``TorchConverter().add_pytorch_module(...).to_coreai()``.
+
+        Raises:
+            ValueError: If ``entrypoint_name`` is already staged.
+            RuntimeError: If ``export_fn`` fails on ``model``.
         """
         if any(e.entrypoint_name == entrypoint_name for e in self._staged):
             raise ValueError(

docs/guides/externalization.ipynb

@@ -100,7 +100,17 @@
     "        return (input * inv_rms).to(input.dtype) * scale\n",
     "\n",
     "\n",
-    "model = RMSNormComposite().eval()\n",
+    "class Model(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.norm = RMSNormComposite()\n",
+    "\n",
+    "    def forward(self, input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:\n",
+    "        return self.norm(input, scale)\n",
+    "\n",
+    "\n",
+    "# The externalized module must be a submodule: the root module is never matched.\n",
+    "model = Model().eval()\n",
     "sample = (torch.randn(10), torch.randn(10))"
    ]
   },
@@ -130,6 +140,8 @@
     "coreai_program = converter.to_coreai()\n",
     "```\n",
     "\n",
+    "Only submodules are externalized, never the root module passed to `add_pytorch_module`, which is why `RMSNormComposite` is wrapped in `Model` above.\n",
+    "\n",
     "Each instance is externalized using its own attribute values, so two instances configured differently are preserved as distinct composite ops.\n"
    ]
   },