Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Add Rngs KeylessInitializers #5017

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
Merged
copybara-service merged 1 commit into from
Oct 10, 2025
Merged

Add Rngs KeylessInitializers #5017

copybara-service merged 1 commit into from
Oct 10, 2025

Conversation

@cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Oct 9, 2025 •
edited
Loading

What does this PR do?

  • Adds initializer constructors from jax.nn.initializers as methods to Rngs and RngsStream that return KeylessInitializers, these follow the Initializer protocol without the leading key argument. Example
rngs = nnx.Rngs(0, param=1)
w = rngs.lecun_normal()((4, 8)) # from default
z = rngs.params.lecun_uniform()((6, 10))
  • Also updates the flax_gspmd guide to use the Rngs initializer plus some other small change such as aliasing PartitionSpec as P and using [...] instead of .value.

@review-notebook-app
Copy link

review-notebook-app bot commented Oct 9, 2025

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

@google-cla
Copy link

google-cla bot commented Oct 9, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@cgarciae cgarciae force-pushed the update-smpd-guide branch 2 times, most recently from b85ff68 to cceda5d Compare October 9, 2025 23:03
@cgarciae cgarciae force-pushed the update-smpd-guide branch 2 times, most recently from c92e2db to 3b7d72b Compare October 10, 2025 00:28
@copybara-service copybara-service bot merged commit b38601d into main Oct 10, 2025
18 checks passed
@copybara-service copybara-service bot deleted the update-smpd-guide branch October 10, 2025 00:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@IvyZX IvyZX IvyZX approved these changes

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@cgarciae @IvyZX