flax
flax copied to clipboard
fix_checkpoint helper
What does this PR do?
This PR moves fix_checkpoint from a code example in the 0.11 migration guide to a helper function.
Fixes # (issue)
Checklist
- [ x] 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 (please add a link).
- [ x] The documentation and docstrings adhere to the documentation guidelines.
- [ ] This change includes necessary high-coverage tests. (No quality testing = no merge!)