# What does this PR do?

<!--

Great, you are contributing to Flax!

But... please read the following carefully so we can make sure your PR is merged
easily.

Replace this text block with a description of the change and which issue it
fixes (if applicable). Please also include relevant motivation/context.

Once you're done, someone in the Flax team will review your PR shortly. They may
suggest changes to make the code even better. If no one reviewed your PR after a
week has passed, don't hesitate to post a new comment @-mentioning the same
persons (sometimes notifications get lost).
-->

Fixes # (issue)

## Checklist
- [ ] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
      checks if that's the case).
- [ ] This change is discussed in a Github issue/
      [discussion](https://github.com/google/flax/discussions) (please add a
      link).
- [ ] The documentation and docstrings adhere to the
      [documentation guidelines](https://github.com/google/flax/blob/main/docs/README.md#how-to-write-code-documentation).
- [ ] This change includes necessary high-coverage tests.
      (No quality testing = no merge!)
