Skip to content
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

Add linen metadata conversion to linx #4552

Merged
merged 1 commit into from
Feb 19, 2025
Merged

Add linen metadata conversion to linx #4552

merged 1 commit into from
Feb 19, 2025
+57 −20

Conversation

IvyZX
Copy link
Collaborator

@IvyZX IvyZX commented Feb 19, 2025 •
edited
Loading

Added Linen variable metadata box conversion to nnx.bridge.Module's init() and apply(). They will now take and output Linen metadata boxes instead of VariableStates.

Added bridge.with_partitioning which uses nnx.with_partitioning within but will automatically converts to nn.Partitioned.

  • For custom metadata class, directly use nnx.with_metadata(..., linen_meta_type=<custom_type>).

@IvyZX IvyZX requested a review from cgarciae February 19, 2025 01:49
cgarciae
cgarciae approved these changes Feb 19, 2025
View reviewed changes
Copy link
Collaborator

@cgarciae cgarciae left a comment

Choose a reason for hiding this comment

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

This is great! Thanks @IvyZX !

@copybara-service copybara-service bot merged commit 9d2bea3 into google:main Feb 19, 2025
16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@cgarciae cgarciae cgarciae 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
@IvyZX @cgarciae