Skip to content

convert RMSNorm to NNX #1728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 1 commit into
base: test_748311465
Choose a base branch
from
Draft

convert RMSNorm to NNX #1728

wants to merge 1 commit into from

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented May 13, 2025 •
edited
Loading

Description

Converts RMSNorm to NNX and adds the rms_norm helper.

Tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

@cgarciae cgarciae changed the base branch from main to test_748311465 May 13, 2025 01:01
bvandermoon
bvandermoon reviewed May 13, 2025
View reviewed changes
Copy link
Collaborator

@bvandermoon bvandermoon left a comment

Choose a reason for hiding this comment

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

Do you know why the full test suite didn't run for the checks?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@bvandermoon bvandermoon bvandermoon left review comments

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk will be requested when the pull request is marked ready for review gobbleturk is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit will be requested when the pull request is marked ready for review khatwanimohit is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla will be requested when the pull request is marked ready for review vipannalla is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan will be requested when the pull request is marked ready for review RissyRan is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 will be requested when the pull request is marked ready for review richjames0 is a code owner

@rni418 rni418 Awaiting requested review from rni418 rni418 will be requested when the pull request is marked ready for review rni418 is a code owner

@gagika gagika Awaiting requested review from gagika gagika will be requested when the pull request is marked ready for review gagika is a code owner

@shralex shralex Awaiting requested review from shralex shralex will be requested when the pull request is marked ready for review shralex is a code owner

@yangyuwei yangyuwei Awaiting requested review from yangyuwei yangyuwei will be requested when the pull request is marked ready for review yangyuwei is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC will be requested when the pull request is marked ready for review SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo will be requested when the pull request is marked ready for review hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha will be requested when the pull request is marked ready for review A9isha is a code owner

@wang2yn84 wang2yn84 Awaiting requested review from wang2yn84 wang2yn84 will be requested when the pull request is marked ready for review wang2yn84 is a code owner

@wyzhang wyzhang Awaiting requested review from wyzhang wyzhang will be requested when the pull request is marked ready for review wyzhang is a code owner

@mitalisi mitalisi Awaiting requested review from mitalisi mitalisi will be requested when the pull request is marked ready for review mitalisi is a code owner

@gpolovets1 gpolovets1 Awaiting requested review from gpolovets1 gpolovets1 will be requested when the pull request is marked ready for review gpolovets1 is a code owner

@mailvijayasingh mailvijayasingh Awaiting requested review from mailvijayasingh mailvijayasingh will be requested when the pull request is marked ready for review mailvijayasingh is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@cgarciae @bvandermoon