Add onnxscript.nn module with Module and Parameter classes (#2819)

This pull request introduces a new PyTorch-like module interface for
building ONNX graphs, enabling users to define reusable neural network
components and manage parameters in a structured way. The main changes
add the `onnxscript.nn` package, expose its API, and implement core
classes for module and parameter management.

Addition of ONNX neural network module interface:

* Added new `onnxscript.nn` package and exposed it in the main
`onnxscript` API, allowing users to access neural network module
functionality. (`onnxscript/__init__.py`,
[[1]](diffhunk://#diff-a1562cf8b37bd59e756ae3802f64fb4b0712845c4dd5a747d7db2a4212f1bfa5R8)
[[2]](diffhunk://#diff-a1562cf8b37bd59e756ae3802f64fb4b0712845c4dd5a747d7db2a4212f1bfa5L130-R131)
* Introduced `Module` and `Parameter` classes in
`onnxscript/nn/_module.py` and `onnxscript/nn/_parameter.py`, providing
a PyTorch-like interface for defining ONNX graph modules, registering
parameters, and managing module hierarchies.
(`onnxscript/nn/_module.py`,
[[1]](diffhunk://#diff-26fd1ef9987845a4613d3d4b2e8d0c6d276b4946ba1a6a3e5e77b6e156e224beR1-R206);
`onnxscript/nn/_parameter.py`,
[[2]](diffhunk://#diff-94ffcabca05fa631e1902fbd4a2e967c00f9db0dca95baae367daef60cd8a7e1R1-R66)
* Created `onnxscript/nn/__init__.py` to expose `Module` and `Parameter`
as the public API of the new package. (`onnxscript/nn/__init__.py`,
[onnxscript/nn/__init__.pyR1-R9](diffhunk://#diff-db47492cd4dbb1a882423614592cb391e31abdec6b5aff85d6f1d2e2716507b4R1-R9))

Core module and parameter functionality:

* `Module` class supports automatic registration of parameters and child
modules, implements methods for iterating over parameters/modules, and
provides `state_dict`/`load_state_dict` for parameter
serialization/deserialization, mirroring PyTorch's API.
(`onnxscript/nn/_module.py`,
[onnxscript/nn/_module.pyR1-R206](diffhunk://#diff-26fd1ef9987845a4613d3d4b2e8d0c6d276b4946ba1a6a3e5e77b6e156e224beR1-R206))
* `Parameter` class subclasses `ir.Value`, allowing direct use in ONNX
ops and supporting initialization, realization, and representation of
parameter tensors. (`onnxscript/nn/_parameter.py`,
[onnxscript/nn/_parameter.pyR1-R66](diffhunk://#diff-94ffcabca05fa631e1902fbd4a2e967c00f9db0dca95baae367daef60cd8a7e1R1-R66))

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

onnxscript/__init__.py

@@ -5,6 +5,7 @@
     "script",
     "graph",
     "ir",
+    "nn",
     "optimizer",
     "rewriter",
     "version_converter",
@@ -127,7 +128,7 @@
 
 # isort: on
 
-from . import ir, optimizer, rewriter, version_converter
+from . import ir, nn, optimizer, rewriter, version_converter
 from ._internal.builder import GraphBuilder
 from ._internal.utils import external_tensor
 from ._internal.values import OnnxFunction, TracedOnnxFunction

onnxscript/_internal/builder.py

@@ -80,10 +80,10 @@ def __init__(self, graph: ir.Graph) -> None:
 
         self._op_builder = self.opset("", opset_version)
 
-        # Context stack to manage hierarchical naming. Each module/layer can push a new context, and pop it when done.
-        # The current context is used as a prefix for naming values and nodes.
-        # This allows us to generate names like "layer1.attention.query"
-        self._context_stack: list[str] = [""]
+        # Module scope stack. Each entry is (name, class_name) where name is
+        # the module attribute name (e.g. "layers.0", "self_attn") and
+        # class_name is the qualified class name (e.g. "Gemma3DecoderLayer").
+        self._scope_stack: list[tuple[str, str]] = []
 
         # Cache for constant initializers (scalars and sequences), keyed by (value, dtype).
         # This avoids creating duplicate initializers for the same constant
@@ -109,7 +109,7 @@ def initializer(
         if name is None:
             name = tensor.name
         if qualify:
-            name = self.qualify_name(name)
+            name = self._qualify_initializer_name(name)
         shape = ir.Shape(tensor.shape)
         value = ir.Value(
             name=name, shape=shape, type=ir.TensorType(tensor.dtype), const_value=tensor
@@ -180,24 +180,26 @@ def _adapt_outputs(
         self, outputs: int | Sequence[str | ir.Value], op_type: str = ""
     ) -> Sequence[ir.Value]:
         if isinstance(outputs, int):
+            count = self.graph.num_nodes()
             if outputs < 0:
                 raise ValueError(f"Number of outputs must be non-negative, got {outputs}")
             if outputs == 1:
-                name = f"{op_type}_output" if op_type else "output"
-                return [ir.Value(name=self.qualify_name(name))]
+                name = f"{op_type}_{count}" if op_type else f"{count}"
+                return [ir.Value(name=self._qualify_value_name(name))]
             else:
                 names = [
-                    f"{op_type}_output{i}" if op_type else f"output{i}" for i in range(outputs)
+                    (f"{op_type}_{count}_{i}" if op_type else f"{count}_{i}")
+                    for i in range(outputs)
                 ]
-                return [ir.Value(name=self.qualify_name(n)) for n in names]
+                return [ir.Value(name=self._qualify_value_name(n)) for n in names]
         adapted_outputs = []
         for output in outputs:
             if isinstance(output, ir.Value):
                 if output.name:
-                    output.name = self.qualify_name(output.name)
+                    output.name = self._qualify_value_name(output.name)
                 adapted_outputs.append(output)
             elif isinstance(output, str):
-                adapted_outputs.append(ir.Value(name=self.qualify_name(output)))
+                adapted_outputs.append(ir.Value(name=self._qualify_value_name(output)))
             else:
                 raise TypeError("Output type not supported.")
         return adapted_outputs
@@ -304,7 +306,7 @@ def call_op(
         outputs = kwargs.pop("_outputs", 1)
 
         count = self.graph.num_nodes()
-        node_name = self.qualify_name(f"{op_type}_node_{count}")
+        node_name = self._qualify_node_name(f"{op_type}_node_{count}")
 
         output_values = self._adapt_outputs(outputs, op_type)
 
@@ -322,33 +324,80 @@ def call_op(
             version=version,
             name=node_name,
         )
+
+        # Attach scope metadata to the node
+        node.metadata_props["namespace"] = self._build_namespace()
+        node.metadata_props["pkg.onnxscript.class_hierarchy"] = repr(self._scope_classes())
+        node.metadata_props["pkg.onnxscript.name_scopes"] = repr(self._scope_names())
+
         self.add_node(node)
 
         return node.outputs if len(node.outputs) > 1 else node.outputs[0]
 
-    def push_module(self, module: str) -> None:
-        """Push a new naming context onto the stack (e.g. a layer or module name)."""
-        current = self.context_name()
-        if module:
-            new_context = f"{current}.{module}" if current else module
-        else:
-            new_context = current
-        self._context_stack.append(new_context)
+    def push_module(self, module: str, class_name: str = "") -> None:
+        """Push a new module scope onto the stack.
+
+        Args:
+            module: The attribute name of the module (e.g. ``"layers.0"``).
+            class_name: The qualified class name (e.g. ``"Gemma3DecoderLayer"``).
+        """
+        self._scope_stack.append((module, class_name))
 
     def pop_module(self) -> None:
-        """Pop the most recent naming context off the stack."""
-        if len(self._context_stack) <= 1:
+        """Pop the most recent module scope off the stack."""
+        if not self._scope_stack:
             raise RuntimeError("Cannot pop_module: no module context has been pushed.")
-        self._context_stack.pop()
+        self._scope_stack.pop()
+
+    def _scope_names(self) -> list[str]:
+        """Return the list of module attribute names in the current scope."""
+        return [name for name, _ in self._scope_stack]
+
+    def _scope_classes(self) -> list[str]:
+        """Return the list of class names in the current scope."""
+        return [cls for _, cls in self._scope_stack]
 
-    def context_name(self) -> str:
-        """Return the current dot-separated naming context prefix."""
-        return self._context_stack[-1] if self._context_stack else ""
+    def _scope_name_parts(self) -> list[str]:
+        """Return non-empty module names for qualifying names."""
+        return [name for name, _ in self._scope_stack if name]
 
-    def qualify_name(self, name: str) -> str:
-        """Prepend the current hierarchical context prefix to the given name."""
-        prefix = self.context_name()
-        return f"{prefix}.{name}" if prefix else name
+    def _qualify_initializer_name(self, name: str) -> str:
+        """Prepend the current hierarchical context prefix to the given name.
+
+        Uses ``.`` as separator, appropriate for parameter and initializer names.
+        """
+        parts = self._scope_name_parts()
+        if parts:
+            return ".".join(parts) + "." + name
+        return name
+
+    def _qualify_value_name(self, name: str) -> str:
+        """Qualify a value name with the current scope using ``.`` separator.
+
+        The name is prefixed with ``v_`` to distinguish values from parameters.
+        """
+        parts = self._scope_name_parts()
+        if parts:
+            return "v_" + ".".join(parts) + "." + name
+        return f"v_{name}"
+
+    def _qualify_node_name(self, name: str) -> str:
+        """Qualify a node name with the current scope using ``/`` separator."""
+        parts = self._scope_name_parts()
+        if parts:
+            return "/".join(parts) + "/" + name
+        return name
+
+    def _build_namespace(self) -> str:
+        """Build the namespace string for a node.
+
+        Each scope entry is formatted as ``name: class_name`` joined by ``/``.
+        """
+        parts = []
+        for name, cls in self._scope_stack:
+            if name or cls:
+                parts.append(f"{name}: {cls}" if cls else name)
+        return "/".join(parts)
 
 
 class OpBuilder: