google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

5.78k stars 610 forks source link

[nnx] add Filters guide #4021

Closed cgarciae closed 1 week ago

cgarciae commented 1 week ago

What does this PR do?

Adds filters_guide to the docs.
Improves current Filters by making them hashable and equatable.

review-notebook-app[bot] commented 1 week ago

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

google-cla[bot] commented 1 week ago

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

IvyZX commented 1 week ago

This is a nice show of the APIs, but I feel it's a bit too minimal and unclear for users to apply it to their use cases.

Some things might worth adding:

How is filter usually used? Maybe have a small model and training loop and show how you can customize trainable params with the API easily.
Various ways to use nnx.split() args. At least showcase when there's no filter predicates, and when there's a function as filter predicate, etc.
Add API reference links to the functions / types mentioned in the doc.

cgarciae commented 1 week ago

Continued in #4028