Add why nnx #4240
Add why nnx #4240
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
cgarciae
force-pushed
the
why-nnx
branch
5 times, most recently
from
277c4ee to
2da1565
Compare
docs_nnx/why.rst
Outdated
| for SPMD annotations, optimizer metadata, etc. | ||
|
|
||
| One choice we made was to use functional (``compact``) semantics for NN programming via the lazy initialization of parameters, | ||
| this made for concise implementation code and aligned our API with Haiku. However, this also meant that the semantics of |
There was a problem hiding this comment.
nit: "concise implementation"
docs_nnx/why.rst
Outdated
| Inspection | ||
| ^^^^^^^^^^ | ||
| The first improvement is that Flax NNX modules are regular Python objects, so you can easily | ||
| construct and inspect them. Because Flax Linen Modules are lazy some attributes are not available |
There was a problem hiding this comment.
nit:
Because Flax Linen Modules are lazy some attributes are not available upon construction and are only accesible at runtime, this makes it hard to inspect and debug.
"Because Flax Linen Modules are lazy, some attributes are not available upon construction and are only accesible at runtime. This makes it hard to inspect and debug."
docs_nnx/why.rst
Outdated
| z = model.encode(x) | ||
| y = model.decoder(z) | ||
|
|
||
| Note that in Linen, calling submodules directly is might not be possible as they are not initialized |
There was a problem hiding this comment.
Note that in Linen, calling submodules directly is might not be possible as they are not initialized so you must construct a new instance, and also the proper parameter structure must be provided.
"Note that in Linen, calling submodules directly is not possible as they are not initialized. So you must construct a new instance and provide proper parameter structure."
docs_nnx/why.rst
Outdated
|
|
||
| State Handling | ||
| ^^^^^^^^^^^^^^ | ||
| One of the areas where Flax Linen is notoriuosly complex is in handling state. When you either use a |
What does this PR do?
Preview