v0.8.5

[cgarciae]

What's Changed

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

New Contributors

• @mcsmart76 made their first contribution in Minor grammar fixes to NNX documentation. #3953
• @jkarwowski made their first contribution in Fixes for batch norm docs #3982
• @sanderland made their first contribution in Fix pickling of exceptions #4002

Full Changelog: v0.8.4...v0.8.5

---

This discussion was created from the release v0.8.5.