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.

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

[nnx] Add interacting with JAX section to README #3573

Merged
merged 1 commit into from
Dec 19, 2023

Conversation

cgarciae
Copy link
Collaborator

What does this PR do?

Improves the NNX README.


While NNX Modules inherently follow reference semantics, they can be easily converted into a pure functional representation that can be used with JAX transformations and other value-based, functional code.

NNX has two very simple APIs to interact with JAX: `split` and `merge`.
Copy link
Collaborator

Choose a reason for hiding this comment

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

I'm wondering if you can add a small comment on how split and merge differs from bind and unbind in Linen and what the advantages are, which could give the users more incentive to try out NNX

Copy link
Collaborator

@chiamp chiamp left a comment

Choose a reason for hiding this comment

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

LGTM, thanks!

@copybara-service copybara-service bot merged commit f1cb011 into main Dec 19, 2023
21 checks passed
@copybara-service copybara-service bot deleted the nnx-improve-readme branch December 19, 2023 21:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants