google
[nnx] Add NNX WeightNorm.
What does this PR do?
This PR addresses #4426 by implementing an NNX WeightNorm class. This brings the requested normalization class to the NNX library – implemented to be consistent with the corresponding module in Linen.
Checklist
- [x] This change is discussed in #4426.
- [x] The documentation and docstrings adhere to the documentation guidelines.
- [x] This change includes necessary high-coverage tests.
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.
![google-cla[bot] avatar](https://avatars.githubusercontent.com/in/42202?v=4 "Published by a pub.dev verified publisher"){kind=small}
To resolve the pytest issues you will have to add some deprecation warning to the list here: https://github.com/google/flax/blob/6af8fcb0bfac1eb078bdf5ba41a3899a71dcaf81/pyproject.toml#L186-L188
If this is not from your PR we can fix it in main first.
To resolve the pytest issues you will have to add some deprecation warning to the list here:
https://github.com/google/flax/blob/6af8fcb0bfac1eb078bdf5ba41a3899a71dcaf81/pyproject.toml#L186-L188
If this is not from your PR we can fix it in
mainfirst.
I believe this would be outside of this PR.
Hello, I think the variable filter is not correctly working? The code does not apply them to extract only the required variables to be normalized, and there is only case where no variable filter is passed.
Hi @nilq - we'd like to finally land this PR on top of all the recent changes to nnx! We just merged support for SpectralNorm (#4623), and your WeightNorm code is quite similar. Do you think it might be possible to factor out some of the shared functionality?
