Skip to content

Update pytree guide to use nnx.capture - #5370

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
samanklesaria:pytree_guide_capture
Apr 13, 2026
Merged

copybara-service[bot] merged 1 commit into
google:mainfrom
samanklesaria:pytree_guide_capture

Conversation

@samanklesaria

Copy link
Copy Markdown
Collaborator

The pytree guide currently demonstrates using the deprecated Module.sow() with nnx.pop pattern. This PR updates it to use nnx.capture.

@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@IvyZX IvyZX left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Just a minor nit. Please add pull-ready label after you've pushed the fix, thanks!

Comment thread docs_nnx/guides/pytree.md

### sow
`sow` receives a `Variable` type, a `name`, and a `value`, and stores it in the `Module` so it can be retrieved at a later time. As the following example shows, NNX APIs such as `nnx.state` or `nnx.pop` are a good way of retrieving the sowed state, however `pop` is recommended because it explicitly removes the temporary state from the Module.
`sow` receives a `Variable` type, a `name`, and a `value`, and stores it in the `Module` so it can be retrieved at a later time. As the following example shows, the `nnx.capture` function allows you to retrieve the sowed state. See the *Extracting Intermediate Values* guide for more details.

@IvyZX IvyZX Apr 13, 2026 •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Did you forgot to put the link in *Extracting Intermediate Values*?

Update: I ended up fixing it while merging.

@copybara-service
copybara-service Bot merged commit 524847f into google:main Apr 13, 2026
21 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