[Relax][Frontend] Add ParameterList and ParameterDict containers

python/tvm/relax/frontend/nn/__init__.py

@@ -19,7 +19,17 @@
 
 # pylint: disable=redefined-builtin
 from . import op, spec
-from .core import Effect, Module, ModuleDict, ModuleList, Object, Parameter, Tensor
+from .core import (
+    Effect,
+    Module,
+    ModuleDict,
+    ModuleList,
+    Object,
+    Parameter,
+    ParameterDict,
+    ParameterList,
+    Tensor,
+)
 from .exporter import add_extern
 from .extern import ExternModule, ObjectModule, SourceModule
 from .modules import (

python/tvm/relax/frontend/nn/core.py

@@ -625,6 +625,63 @@ def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
             module.to(dtype=dtype)
 
 
+class ParameterDict(Module):
+    """Holds parameters in a dict."""
+
+    def __init__(
+        self,
+        params: OrderedDict[str, Parameter] | dict[str, Parameter] | None = None,
+    ):
+        self.params: OrderedDict[str, Parameter] = OrderedDict()
+        if params is not None:
+            self.update(params)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.params)
+
+    def __getitem__(self, key: str) -> Parameter:
+        return self.params[key]
+
+    def __setitem__(self, key: str, param: Parameter) -> None:
+        if not isinstance(key, str):
+            raise TypeError(f"ParameterDict keys must be strings, but got {type(key).__name__}")
+        if not isinstance(param, Parameter):
+            raise TypeError(f"ParameterDict values must be nn.Parameter, but got {type(param).__name__}")
+        self.params[key] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def keys(self) -> Iterator[str]:
+        return self.params.keys()
+
+    def values(self) -> Iterator[Parameter]:
+        return self.params.values()
+
+    def items(self) -> Iterator[tuple[str, Parameter]]:
+        return self.params.items()
+
+    def get(self, key: str, default: Parameter | None = None) -> Parameter | None:
+        return self.params.get(key, default)
+
+    def update(self, params: dict[str, Parameter]) -> None:
+        for key, param in params.items():
+            self[key] = param
+
+    def clear(self) -> None:
+        self.params.clear()
+
+    def pop(self, key: str) -> Parameter:
+        return self.params.pop(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.params
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params.values():
+            param.to(dtype=dtype)
+
+
 class ModuleList(Module):
     """Holds submodules in a list."""
 
@@ -658,6 +715,44 @@ def forward(self, x):  # pylint: disable=invalid-name
         return x
 
 
+class ParameterList(Module):
+    """Holds parameters in a list."""
+
+    def __init__(self, params: list[Parameter] | None = None):
+        self.params: list[Parameter] = []
+        if params is not None:
+            self.extend(params)
+
+    def __iter__(self) -> Iterator[Parameter]:
+        return iter(self.params)
+
+    def __getitem__(self, idx: int) -> Parameter:
+        return self.params[idx]
+
+    def __setitem__(self, idx: int, param: Parameter) -> None:
+        if not isinstance(param, Parameter):
+            raise TypeError(f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}")
+        self.params[idx] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def append(self, param: Parameter) -> None:
+        """Add a parameter to the end of the ParameterList"""
+        if not isinstance(param, Parameter):  
+            raise TypeError(f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}")         
+        self.params.append(param)
+
+    def extend(self, params: list[Parameter]) -> None:
+        """Add parameters to the end of the ParameterList"""
+        for param in params:
+            self.append(param)
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params:
+            param.to(dtype=dtype)
+
+
 def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
     """Wrap the given relax.Expr, emit it using the current BlockBuilder,
     and automatically handle nested cases if the expr represents a Tuple.
@@ -692,7 +787,17 @@ def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
 
 def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):
     """Find attributes that satisfy the condition recursively"""
-    if isinstance(root, ModuleList):
+    if isinstance(root, ParameterList):
+        for i, param in enumerate(root):
+            if condition_yield(param):
+                yield prefix + f"{i}", param
+        return
+    elif isinstance(root, ParameterDict):
+        for name, param in root.items():
+            if condition_yield(param):
+                yield prefix + name, param
+        return
+    elif isinstance(root, ModuleList):
         for i, subitem in enumerate(root):
             yield from _attribute_finder(subitem, prefix + f"{i}.", condition_yield)
         return
@@ -703,6 +808,18 @@ def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any]
     for name, item in root.__dict__.items():
         if condition_yield(item):
             yield prefix + name, item
+        elif isinstance(item, ParameterList):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ParameterDict):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
         elif isinstance(item, ModuleList):
             yield from _attribute_finder(
                 item,

python/tvm/relax/frontend/nn/visitor.py

@@ -116,6 +116,42 @@ def visit_modulelist(self, name: str, node: nn.ModuleList) -> Any:
         """
         return self.visit(name, node)
 
+    def visit_parameterdict(self, name: str, node: nn.ParameterDict) -> Any:
+        """The base visiting method for mutation of nn.ParameterDict nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterDict
+            The current node of nn.ParameterDict to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_parameterlist(self, name: str, node: nn.ParameterList) -> Any:
+        """The base visiting method for mutation of nn.ParameterList nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterList
+            The current node of nn.ParameterList to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
     def visit(self, name: str, node: Any) -> Any:
         """The base dispatching method for visiting of all nodes.
 
@@ -141,9 +177,19 @@ def _get_child_name(parent: str, child: str) -> str:
             else:
                 return f"{parent}.{child}"
 
-        if isinstance(node, nn.ModuleList):
+        if isinstance(node, nn.ParameterList):
+            for i in range(len(node)):
+                node[i] = self.visit_param(_get_child_name(name, str(i)), node[i])
+        elif isinstance(node, nn.ParameterDict):
+            for k, v in node.items():
+                node[k] = self.visit_param(_get_child_name(name, k), v)
+        elif isinstance(node, nn.ModuleList):
             for i in range(len(node)):
-                if isinstance(node[i], nn.ModuleDict):
+                if isinstance(node[i], nn.ParameterDict):
+                    node[i] = self.visit_parameterdict(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ParameterList):
+                    node[i] = self.visit_parameterlist(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ModuleDict):
                     node[i] = self.visit_moduledict(f"{name}.{i}", node[i])
                 elif isinstance(node[i], nn.ModuleList):
                     node[i] = self.visit_modulelist(f"{name}.{i}", node[i])
@@ -155,7 +201,11 @@ def _get_child_name(parent: str, child: str) -> str:
                     node[i] = self.visit_param(f"{name}.{i}", node[i])
         elif isinstance(node, nn.ModuleDict):
             for k, v in node.items():
-                if isinstance(v, nn.ModuleDict):
+                if isinstance(v, nn.ParameterDict):
+                    node[k] = self.visit_parameterdict(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ParameterList):
+                    node[k] = self.visit_parameterlist(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ModuleDict):
                     node[k] = self.visit_moduledict(_get_child_name(name, k), v)
                 elif isinstance(v, nn.ModuleList):
                     node[k] = self.visit_modulelist(_get_child_name(name, k), v)
@@ -167,7 +217,11 @@ def _get_child_name(parent: str, child: str) -> str:
                     node[k] = self.visit_param(_get_child_name(name, k), v)
         else:
             for key, value in node.__dict__.items():
-                if isinstance(value, nn.ModuleDict):
+                if isinstance(value, nn.ParameterDict):
+                    setattr(node, key, self.visit_parameterdict(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ParameterList):
+                    setattr(node, key, self.visit_parameterlist(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ModuleDict):
                     setattr(node, key, self.visit_moduledict(_get_child_name(name, key), value))
                 elif isinstance(value, nn.ModuleList):
                     setattr(node, key, self.visit_modulelist(_get_child_name(name, key), value))

tests/python/relax/test_frontend_nn_mutator.py

@@ -127,6 +127,42 @@ def visit_param(self, name: str, node: nn.Parameter) -> Any:
     mutator.visit("mod_list", mod_list)
 
 
+def test_mutator_naming_parameter_containers():
+    class Module(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.param_list = nn.ParameterList(
+                [
+                    nn.Parameter((32, 128), "float64"),
+                    nn.Parameter((32, 128), "float32"),
+                ]
+            )
+            self.param_dict = nn.ParameterDict(
+                {
+                    "k0": nn.Parameter((32, 128), "float16"),
+                    "k1": nn.Parameter((32, 128), "float8"),
+                }
+            )
+
+    seen = []
+
+    class Mutator(nn.Mutator):
+        def visit_param(self, name: str, node: nn.Parameter) -> Any:
+            seen.append((name, node.dtype))
+            return node
+
+    module = Module()
+    mutator = Mutator()
+    mutator.visit("", module)
+
+    assert seen == [
+        ("param_list.0", "float64"),
+        ("param_list.1", "float32"),
+        ("param_dict.k0", "float16"),
+        ("param_dict.k1", "float8"),
+    ]
+
+
 def test_mutator_module():
     class SubModule1(nn.Module):
         def __init__(self) -> None: