Upgrade Flax NNX Bridge API doc #4198
Open
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
8bitmp3
force-pushed
the
upgrade-nnx-bridge
branch
from
45476c9
to
5d31452
Compare
8bitmp3
force-pushed
the
upgrade-nnx-bridge
branch
from
28eff8e
to
b3e65f8
Compare
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
|
|
||
| Since Linen modules may require an input to create variables, we semi-formally supported lazy initialization in the NNX modules converted from Linen. The Linen variables are created when you give it a sample input. | ||
| Since `flax.linen.Module`s may require an input to create variables, the Flax team semi-formally supports lazy initialization in the `flax.nnx.Module`s converted from Flax Linen. The Flax Linen variables are created when you give it a sample input. For you, it's calling [`nnx.bridge.lazy_init`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/bridge.html#flax.nnx.bridge.ToNNX.lazy_init) (`nnx.bridge.ToNNX.lazy_init`) where you call `module.init()` in the Flax Linen code. |
There was a problem hiding this comment.
@cgarciae We should mention the ToNNX bridge here since it's used in the code snippet below.
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
|
|
||
| ### NNX to Linen | ||
|
|
||
| If you are not using any metadata feature of the `nnx.Variable` (i.e., no sharding annotation, no registered hooks), the converted Linen module will not add a metadata wrapper to your NNX variable, and you don't need to worry about it. |
There was a problem hiding this comment.
Since this part was rewritten, will revert to the new version too @cgarciae
There was a problem hiding this comment.
If you are not using any metadata features of the [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable) - that is, no sharding annotation, no registered hooks, then the converted `flax.linen.Module` will not add a metadata wrapper to your Flax NNX variable, and you won't need to worry about it. (Recall that all Flax NNX variables are wrapped with `nnx.Variable` box.
But if you did add sharding annotations to your Flax NNX variables, then [`nnx.bridge.ToLinen`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/bridge.html#flax.nnx.bridge.ToLinen) will convert them to a default Flax Linen partition metadata class called [`flax.nnx.bridge.NNXMeta`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/bridge.html#flax.nnx.bridge.NNXMeta), retaining all the metadata you put into the NNX variable.
Similar to any Flax Linen metadata wrappers, you can use `flax.linen.unbox()` ([`flax.linen.meta.unbox`](https://github.com/google/flax/blob/5d31452889b8d106d7c722b5eaac14cb9784fec2/flax/core/meta.py#L160)) to get the raw [`jax.Array`](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) [pytree](https://jax.readthedocs.io/en/latest/glossary.html#term-pytree).
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
8bitmp3
commented
Oct 7, 2024
8bitmp3
force-pushed
the
upgrade-nnx-bridge
branch
2 times, most recently
from
b62ae11
to
2f6e944
Compare
8bitmp3
force-pushed
the
upgrade-nnx-bridge
branch
from
2f6e944
to
b1ba25f
Compare
|
TODO Fix the imports section, fix Markdown in the basic example, fix certain URL rendering. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.