Skip to content

[Relax][Frontend] Add ParameterList and ParameterDict containers#19495

Merged
tlopex merged 2 commits intoapache:mainfrom
mshr-h:add-parameter-list-dict
May 2, 2026
Merged

[Relax][Frontend] Add ParameterList and ParameterDict containers#19495
tlopex merged 2 commits intoapache:mainfrom
mshr-h:add-parameter-list-dict

Conversation

@mshr-h
Copy link
Copy Markdown
Contributor

@mshr-h mshr-h commented May 2, 2026

This PR adds first-class nn.ParameterList and nn.ParameterDict containers to the Relax frontend.

These containers provide PyTorch-like list/dict registration for raw nn.Parameter objects while preserving Relax frontend semantics: values must be explicit nn.Parameter instances, with no automatic tensor-to-parameter conversion.

Changes

  • Add public nn.ParameterList and nn.ParameterDict exports.
  • Support stable parameter names in traversal:
    • params.0, params.1
    • params.foo, params.bar
  • Integrate the new containers with:
    • named_parameters()
    • parameters()
    • state_dict()
    • load_state_dict()
    • to(dtype=...)
    • export_tvm()
    • nn.Mutator
  • Add focused tests for basic container behavior, type validation, nested traversal, export parameter names, state loading, dtype conversion, and mutator naming.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces ParameterDict and ParameterList containers to the Relax frontend's nn module, allowing for structured management of parameters. It includes updates to the Mutator for proper traversal and naming of parameters within these containers, along with comprehensive unit tests. Feedback focuses on adding explicit runtime type validation for keys and values in these containers to ensure consistency with the provided tests and improving naming consistency within the visitor logic.

Comment thread python/tvm/relax/frontend/nn/core.py
Comment thread python/tvm/relax/frontend/nn/core.py
Comment thread python/tvm/relax/frontend/nn/core.py
Comment thread python/tvm/relax/frontend/nn/visitor.py
mshr-h and others added 2 commits May 2, 2026 22:09
Add first-class nn.ParameterList and nn.ParameterDict containers for
Relax frontend parameters. The new containers validate explicit
nn.Parameter values, participate in parameter/state traversal, export
stable dot-separated names, and are handled by nn.Mutator traversal.

Add tests for public exports, list/dict behavior, nested traversal,
state_dict/load_state_dict, dtype conversion, export parameter names,
and mutator naming.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@mshr-h mshr-h force-pushed the add-parameter-list-dict branch from c1ab8de to ee3e9ea Compare May 2, 2026 13:09
@mshr-h mshr-h marked this pull request as ready for review May 2, 2026 13:11
@tlopex tlopex merged commit 86794e7 into apache:main May 2, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants