Update flax requirement from >=0.7.0 to >=0.8.5

[dependabot]

Updates the requirements on flax to permit the latest version.

Release notes

Sourced from flax's releases.

v0.8.5

What's Changed

• v0.8.5 by @​cgarciae in google/flax#3941
• [nnx] improve vmap axis size detection by @​cgarciae in google/flax#3947
• Add direct penzai.treescope support for NNX objects. by @​copybara-service in google/flax#3948
• [nnx] fix nnx_basics dependencies by @​cgarciae in google/flax#3942
• Rename all the NNX tests to internal naming & build conventions. by @​copybara-service in google/flax#3952
• updated rng guide by @​chiamp in google/flax#3912
• upgraded haiku guide to include NNX by @​chiamp in google/flax#3923
• parameterized NNX transforms tests by @​chiamp in google/flax#3906
• Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes. by @​copybara-service in google/flax#3957
• fix HEAD by @​chiamp in google/flax#3960
• Minor grammar fixes to NNX documentation. by @​mcsmart76 in google/flax#3953
• Make FlatState a Mapping instead of a dict by @​NeilGirdhar in google/flax#3928
• Adding Welford metric. by @​copybara-service in google/flax#3959
• Modify Welford metric to return mean value. by @​copybara-service in google/flax#3970
• [nnx] make State generic by @​cgarciae in google/flax#3964
• updated NNX nn docstrings by @​chiamp in google/flax#3972
• make flax work with upcoming JAX change to tree_map (being more careful about by @​copybara-service in google/flax#3976
• updated nnx.module docstrings by @​chiamp in google/flax#3966
• updated nnx.Conv and nnx.ConvTranspose by @​chiamp in google/flax#3974
• updated nnx.graph docstrings by @​chiamp in google/flax#3958
• • Adds pmap and Pmap. static_broadcasted_argnums, donate_argnums, and global_arg_shapes are not yet supported. by @​copybara-service in google/flax#3978
• Fixes for batch norm docs by @​jkarwowski in google/flax#3982
• fix deprecation warning by @​chiamp in google/flax#3981
• updated NNX rnglib docstring by @​chiamp in google/flax#3980
• updated nnx.training by @​chiamp in google/flax#3975
• updated nnx.variables docstrings by @​chiamp in google/flax#3986
• [nnx] vectorize vmap split counts by @​cgarciae in google/flax#3989
• added wrt option to nnx.Optimizer by @​chiamp in google/flax#3983
• Added nnx.graph.iter_children by @​chiamp in google/flax#3991
• [nnx] fix vmap by @​copybara-service in google/flax#3995
• Fix head pytest breakage by @​IvyZX in google/flax#4006
• Helper function for loading params from a linen module by @​copybara-service in google/flax#4012
• Port gemma/layers to NNX by @​copybara-service in google/flax#4013
• [nnx] fix grad by @​cgarciae in google/flax#4007
• [nnx] add PathContains Filter by @​cgarciae in google/flax#4011
• Support Python 3.9 by @​copybara-service in google/flax#4018
• Port gemma/modules to NNX by @​copybara-service in google/flax#4014
• Internal change to fix current head CI by @​copybara-service in google/flax#4017
• Unpin the Orbax pip version. by @​copybara-service in google/flax#4024
• Fix Gemma test to unbreak head by @​IvyZX in google/flax#4025
• Fix pickling of exceptions by @​sanderland in google/flax#4002
• Call user-defined variable transforms before determining axis size in nn.vmap. by @​copybara-service in google/flax#4026
• CI: add test run against oldest supported jax version by @​jakevdp in google/flax#3996
• Make force_fp32_for_softmax arg in MultiHeadDotProductAttention useful. by @​copybara-service in google/flax#4029

New Contributors

• @​mcsmart76 made their first contribution in google/flax#3953
• @​jkarwowski made their first contribution in google/flax#3982

... (truncated)

Changelog

Sourced from flax's changelog.

Changelog

vNext

(Add your change to a random empty line to avoid merge conflicts)

• removed GeGLU simplistic activation, it should be implemented manually.
• removed FLAX_LAZY_RNG flag support for old non-lazy PRNG derivation mode

0.8.2

• fixed rng guide outputs by @​chiamp in google/flax#3685
• enforce mask kwarg in norm layers by @​chiamp in google/flax#3663
• added kwargs to self.param and self.variable by @​chiamp in google/flax#3675
• added nnx normalization tests by @​chiamp in google/flax#3689
• added NNX init_cache docstring example by @​chiamp in google/flax#3688
• added nnx attention equivalence test by @​chiamp in google/flax#3687
• Fix bug that assumed frozen-dict keys were strings. by @​copybara-service in google/flax#3692
• added nnx rmsnorm by @​chiamp in google/flax#3691
• updated nnx compute_stats by @​chiamp in google/flax#3693
• fixed intercept_methods docstring by @​chiamp in google/flax#3694
• [nnx] Add Sphinx Docs by @​cgarciae in google/flax#3678
• Fix pointless docstring example of nn.checkpoint / nn.remat. by @​levskaya in google/flax#3703
• added default params rng to .apply by @​chiamp in google/flax#3698
• [nnx] add partial_init by @​cgarciae in google/flax#3674
• make make_rng default to 'params' by @​chiamp in google/flax#3699
• Add SimpleCell. by @​carlosgmartin in google/flax#3697
• fix Module.module_paths docstring by @​cgarciae in google/flax#3709
• Guarantee the latest JAX version on CI by @​cgarciae in google/flax#3705
• Replace deprecated API jax.tree.map by @​copybara-service in google/flax#3715
• Use jax.tree_util.tree_map instead of deprecated jax.tree.map. by @​copybara-service in google/flax#3714
• [nnx] simplify readme by @​cgarciae in google/flax#3707
• [nnx] add demo.ipynb by @​cgarciae in google/flax#3680

... (truncated)

Commits

• 9e6b819 Make force_fp32_for_softmax arg in MultiHeadDotProductAttention useful.
• 3b21870 Merge pull request #3996 from jakevdp:oldest-jax
• e5cb2f7 Call user-defined variable transforms before determining axis size in nn.vmap.
• 58020c4 CI: add test run against oldest supported jax version
• 791df52 Merge pull request #4002 from sanderland:fixpickle
• 9b6748a Merge pull request #4025 from IvyZX:unbreak
• 1c4e348 Fix Gemma test to unbreak head
• 68d5c81 Unpin the Orbax pip version.
• bda7da1 Internal change to fix current head CI
• 9ca8db5 Port gemma/modules to NNX
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)