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

Add NNX checkpointing guide #4249

Merged
merged 1 commit into from
Oct 7, 2024
Merged

Conversation

IvyZX
Copy link
Collaborator

@IvyZX IvyZX commented Oct 4, 2024

  • Made a simple checkpointing guide for NNX, focusing on showcasing Flax side of APIs and referring to Orbax site for Orbax side API
  • Fixed the surgery guide a bit, simplified its code
  • Added a simple config check for nnx.Linear
  • Made sure pre-commit don't check markdown files for docs_nnx either, otherwise we'll be in a whirlwind of whitespace reads-and-writes

@IvyZX IvyZX requested a review from cgarciae October 4, 2024 02:03
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@@ -0,0 +1,201 @@
# Save and load checkpoints
Copy link
Collaborator

Choose a reason for hiding this comment

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

missing entry for the document in index.rst

@copybara-service copybara-service bot merged commit 9a3a1fc into google:main Oct 7, 2024
18 checks passed
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