**Issue description:** **Steps to reproduce:** **Minimal reproduction project:** **Flax version:**