[nnx] add Using Filters guide #4028
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
Hey @IvyZX, something weird happened with the other PR so I am continuing it here. Thanks for reviewing the PR, sadly I hadn't pushed all my changes and the guide was half way through. Check out the new content. I focused the guide on making it clear how filters work and how state is filtered. Let me know if you still think I should add more examples. |
|
@cgarciae this is beautiful! Makes a big difference in forming a mental model on NNX. Is a new release planned anytime the next month or so which encapsulates this sort of stuff? I would love to prepare some of these notes fall teaching to show the NNX workflow |
docs/nnx/filters_guide.md
Outdated
| | | `PathContains(key)` | Matches values that have an associated `path` that contains the given `key` | | ||
| | `'{filter}'` | `WithTag('{filter}')` | Matches values that have string `tag` attribute equal to `'{filter}'` | | ||
| | `(*filters)` | `Any(*filters)` | Matches values that match any of the inner `filters` | | ||
| | `[*filters]` | `Any(*filters)` | Matches values that match all of the inner `filters` | |
There was a problem hiding this comment.
Is some part of this grid duplicated? Not sure I understood it fully but the tuple is "any" while the list is "all"? This feels a bit confusing...
There was a problem hiding this comment.
Merged them into 1 row to clarify that both tuples and lists are converted to Any
| Reversing the order will make sure that the `SpecialParam` are captured first | ||
|
|
||
| ```{code-cell} | ||
| graphdef, special_params, params = split(bar, SpecialParam, nnx.Param) # correct! |
There was a problem hiding this comment.
Since we introduced the filter DSL above, it would be cool to showcase it here in this example!
There was a problem hiding this comment.
Added an example using the DSL in nnx.vmap
|
@jlperla coincidentally I just created a new release (needed it for a bug fix), I can create another one soon. However, all the features shown here are stable so it should not impact the information in the guide a lot (except for some types being exposed at the top-level). |
|
Thanks @cgarciae If the interfaces are stable then that is good for me to experiment. I will try some of this out in the next month and then update my lecture notes august (at which point I can point to new docs) so no time pressure on my side. Looking forward to trying out the library. Also keen to see in august how the "training loop" pattern with metrics, optax, etc. that I can teach as a boilerplate - which I noticed you were working on. Lots of movement! |
cgarciae
force-pushed
the
nnx-filters-guide2
branch
2 times, most recently
from
36f2b32
to
2ae48a7
Compare
cgarciae
force-pushed
the
nnx-filters-guide2
branch
from
2ae48a7
to
0aa5fe6
Compare
What does this PR do?
filters_guideto the docs: preview