math/py-flax: Update to 0.6.4

Changes: https://github.com/google/flax/releases