Upgrade Flax NNX Bridge API doc #4198

8bitmp3 · 2024-09-16T20:55:13Z

No description provided.

review-notebook-app · 2024-09-16T20:55:18Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

docs_nnx/guides/bridge_guide.md

8bitmp3 · 2024-10-07T11:47:19Z

docs_nnx/guides/bridge_guide.md


-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.


@cgarciae We should mention the ToNNX bridge here since it's used in the code snippet below.

docs_nnx/guides/bridge_guide.md

8bitmp3 · 2024-10-07T12:01:39Z

docs_nnx/guides/bridge_guide.md


-### 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.


Since this part was rewritten, will revert to the new version too @cgarciae

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).

docs_nnx/guides/bridge_guide.md

8bitmp3 · 2024-10-08T12:53:55Z

TODO Fix the imports section, fix Markdown in the basic example, fix certain URL rendering.

8bitmp3 closed this Oct 6, 2024

8bitmp3 force-pushed the upgrade-nnx-bridge branch from 45476c9 to 5d31452 Compare October 6, 2024 21:53

8bitmp3 reopened this Oct 7, 2024

8bitmp3 requested review from IvyZX and cgarciae October 7, 2024 00:30

8bitmp3 self-assigned this Oct 7, 2024

8bitmp3 force-pushed the upgrade-nnx-bridge branch from 28eff8e to b3e65f8 Compare October 7, 2024 00:47