Version 0.7.3

What's Changed

• Fix Documentation Typo by @peterdavidfagan in #3265
• bump version number by @chiamp in #3273
• added cursor api by @chiamp in #3246
• Improve Typing support by @cgarciae in #3242
• feat: add configs for vae example by @SauravMaheshkar in #3254
• fix stackoverflow when loading pickled module by @cgarciae in #3286
• Remove string from checkpointing example and unbreak doctest in head by @IvyZX in #3288
• Improve kw_only_dataclass by @cgarciae in #3293
• updated haiku upgrade guide by @chiamp in #3292
• added logical partitioning to pjit guide by @chiamp in #3290
• Rolling back ba9e24a by @copybara-service in #3295
• Add RNN FLIP by @cgarciae in #2585
• [flax] Add module path to nn.module. by @copybara-service in #3300
• Minor fix to jax_utils.prefect_to_device by @anuragarnab in #3308
• add tf-text to doc requirements by @cgarciae in #3317
• Allow apply's method argument to accept submodules by @cgarciae in #3281
• Update handling of typed PRNG keys by @jakevdp in #3314
• 0.7.3 release by @IvyZX in #3326

New Contributors

• @peterdavidfagan made their first contribution in #3265
• @SauravMaheshkar made their first contribution in #3254

Full Changelog: v0.7.2...v0.7.3

Version 0.7.2

What's Changed

• make flax.core.copy add_or_replace optional by @PhilipVinc in #3241
• bump version to 0.7.2 by @chiamp in #3258
• Release 0.7.2 by @chiamp in #3269

Full Changelog: v0.7.1...v0.7.2

Version 0.7.1

What's Changed

• added dict migration guide to index by @chiamp in #3188
• [linen] Add alternative, more numerically stable, variance calculation to LayerNorm. by @copybara-service in #3194
• [linen] Minor cleanup to normalization code. by @copybara-service in #3200
• Fix warnings from atari gym. by @levskaya in #3207
• Fix carry slice logic by @cgarciae in #3213
• removed FrozenDict section by @chiamp in #3210
• make flax_basics guide use utility fns by @chiamp in #3214
• Fix typo in struct.py documentation by @marcoselvi in #3209
• updated flax guides to use utility fns by @chiamp in #3215
• [linen] More minor cleanup in normalization compute_stats. by @copybara-service in #3205
• Expose options to customize rich.Table by @cgarciae in #3197
• Use pyink by @cgarciae in #3216
• Fix remaining pyink issues by @cgarciae in #3219
• add scan over layers section by @cgarciae in #3195
• Fix checkpointing guide error at head by @IvyZX in #3223
• Improve scan docs by @cgarciae in #3231
• 0.7.1 release by @chiamp in #3234

New Contributors

• @marcoselvi made their first contribution in #3209

Full Changelog: v0.7.0...v0.7.1

Version 0.7.0

What's Changed

• delete long-unused dotgetter utility by @copybara-service in #3156
• Add sections to flax.linen toctree by @cgarciae in #3073
• update doctest requirements to use recentish jax minver by @levskaya in #3160
• Update gym to gymnasium by @cgarciae in #3133
• Add lifted transforms section to Haiku migration guide by @cgarciae in #3158
• relaxed tolerance by @chiamp in #3162
• fixed broken links by @chiamp in #3161
• added dict migration guide by @chiamp in #3109
• Update requirements, restructure files and fix formatting for VAE example by @canyon289 in #3046
• Update python version support by @cgarciae in #3168
• Improve Haiku migration guide by @cgarciae in #3169
• Set default types in Flax for Orbax restoration and add restore_with_serialized_types in preparation for an upcoming change. by @copybara-service in #3165
• fix precommit issues by @chiamp in #3170
• fixed incorrect reference link by @chiamp in #3167
• added absltest to linen_recurrent_test by @chiamp in #3172
• remove cell_size docs from RNN by @cgarciae in #3186
• 0.7.0 by @cgarciae in #3187

Full Changelog: v0.6.11...v0.7.0

Version 0.6.11

What's Changed

• remove unreferenced docs by @chiamp in #3085
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/ogbg_molpcba by @dependabot in #2976
• Bind submodules using getattribute by @cgarciae in #2901
• Allow manually specifying the rng key for Dropout by @cgarciae in #3114
• fixed missing separator for rng fold in by @chiamp in #3095
• Bump imagenet Flax to 0.6.5 by @IvyZX in #3118
• Rolling back #3104 because internal tests are breaking. by @copybara-service in #3126
• Use pyproject.toml as hash on CI by @cgarciae in #3131
• Minor improvements for Sequential by @cgarciae in #3129
• Update activation.py docstring for PReLU by @levskaya in #3122
• Eagerly bind submodules by @cgarciae in #3077
• Add haiku_migration_guide by @cgarciae in #3139
• Update Haiku migration guide. by @LenaMartens in #3154
• Fix setup trigger issue under sharing and transforms. by @levskaya in #3153
• Upgrade the pjit guide to jax.jit by @IvyZX in #3016
• Release 0.6.11 by @IvyZX in #3155

New Contributors

• @LenaMartens made their first contribution in #3154

Full Changelog: v0.6.10...v0.6.11

Version 0.6.10

What's Changed

• Add missing Bidirectional api reference by @cgarciae in #2999
• Optimize struct.dataclass by @cgarciae in #3024
• Rectified typos & grammar in docs by @Neilblaze in #3021
• Updated full_eval.ipynb & corrected to constructor by @Neilblaze in #3059
• fixed broken coverage link by @chiamp in #3058
• Allow clipping the dynamic loss scale to a minimum value by @copybara-service in #3051
• updated extracting_intermediates doc by @chiamp in #3063
• updated api reference by @chiamp in #3068
• fix RTD urllib3 error by @chiamp in #3081
• Add test case for degenerate and deterministic attention layer dropout. by @copybara-service in #3083
• 0.6.10 by @cgarciae in #3089

New Contributors

• @Neilblaze made their first contribution in #3021

Full Changelog: v0.6.9...v0.6.10

Version 0.6.9

What's Changed

• Follow-up to #2985 (adding config.prefetch) by @copybara-service in #3004
• Bringing changes in #2955 internally. by @copybara-service in #2960
• Updates imagenet.ipynb to use GPU Colab runtime. by @andsteing in #3007
• Simplify pyproject.toml's mypy section by @cgarciae in #3011
• Invalidate CI cache daily by @cgarciae in #3010
• Update module.py by @aniquetahir in #3006
• Fix get_partition_spec on replicated array by @IvyZX in #3013
• Move setup.py to project.toml by @cgarciae in #3012
• modify git repo master to main by @Mddct in #3020
• Added pretty_repr utility fn by @chiamp in #3022
• Update checkpointing guide with Orbax APIs by @IvyZX in #3002
• Checkpoint fixes to unbreak head by @IvyZX in #3027
• Use backend dependent os.path.getsize. by @copybara-service in #3028
• Add more explicit warnings and logs for code with Orbax backend. by @copybara-service in #3030
• Update orbax_upgrade_guide.rst fixing broken link by @canyon289 in #3031
• Back to orbax-checkpoint after the package is fixed in 0.2.1. by @IvyZX in #3029
• 0.6.9 release by @IvyZX in #3035
• Fix release CI script by @cgarciae in #3039

New Contributors

• @aniquetahir made their first contribution in #3006
• @Mddct made their first contribution in #3020
• @canyon289 made their first contribution in #3031

Full Changelog: v0.6.8...v0.6.9

Version 0.6.8

Changelog

• The automatic checkpoint migration was temporarily rolled back due to legacy compatibility issues.
  • We still recommend you to use the upgrade guide and migrate completely to the Orbax API to ensure stability.
  • Or alternatively, add flax.config.update('flax_use_orbax_checkpointing', True) to your project to avoid being impacted by the automatic migration process.
• Added utility functions to frozen_dict api.
• Migrated Flax away from register_keypaths.
• Fixes kwargs in convert_to_graphs_tuple_fn.
• Fixed examples in a few ways:
  • Bumped the TF version
  • Used latest checkpoint formats
  • Other misc fixes.

List of PRs

• Add Google Research examples with Flax by @8bitmp3 in #2934
• [Trivial] [Post-release] Bump up the HEAD version number by @IvyZX in #2946
• fix tabulate duplicating parameters by @cgarciae in #2931
• permit passing mesh or rules to spmd primitives by @cgarciae in #2939
• fixed ScopeParamShapeError message by @chiamp in #2947
• Fix checkpointing docs for Orbax 0.1.4 by @IvyZX in #2962
• Migrate Flax away from register_keypaths by @copybara-service in #2963
• set axis='batch' by @yechengxi in #2967
• Add Bidirectional by @cgarciae in #2770
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/ppo by @dependabot in #2984
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/imagenet by @dependabot in #2983
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/sst2 by @dependabot in #2980
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/nlp_seq by @dependabot in #2979
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/lm1b by @dependabot in #2977
• Adds missing api references by @cgarciae in #2926
• Added utility functions to frozen_dict api by @chiamp in #2975
• Fixed pop utility fn doc string by @chiamp in #2958
• release 0.6.8 by @IvyZX in #2987

Full Changelog: v0.6.7...v0.6.8

Version 0.6.7

What's Changed

• Fix lm1b sampler by @cgarciae in #2842
• RNN Module by @cgarciae in #2604
• Add Orbax migration guide by @IvyZX in #2885
• Add 'all' install group by @cgarciae in #2894
• Improve nn.scan docs by @cgarciae in #2839
• make utility functions work on FrozenDicts and regular dicts by @chiamp in #2919
• Improve Flax dark mode - remove hard-coded text color by @8bitmp3 in #2924
• Fix pytype and improve typing for Module by @cgarciae in #2933
• Make flax.version 1 + the current release by @IvyZX in #2936
• Migrate Flax dataclass to the newest JAX pytree keypath API. by @copybara-service in #2921
• release 0.6.7 by @IvyZX in #2944

Full Changelog: v0.6.6...v0.6.7

Version 0.6.6

What's Changed

• Enhance checkpointing docs with Orbax information by @cpgaffney1 in #2879
• A small doc and one fixe to Imagenet benchmark by @nouiz in #2812
• Refactor and separate out name relaxation policy changes. by @levskaya in #2873
• Added RMS normalization layer by @chiamp in #2881
• added feature flag for FrozenDicts by @chiamp in #2872
• Explicitly mark the lazy-rng with legacy_config, and all new configs should use flax.config. by @copybara-service in #2895
• Added test for FrozenDict feature flag by @chiamp in #2896
• Added dict utility functions by @chiamp in #2897
• Improve notebook handling by @cgarciae in #2860
• migrate tests to regular dicts by @chiamp in #2898
• Release 0.6.6 by @IvyZX in #2908

New Contributors

• @cpgaffney1 made their first contribution in #2879

Full Changelog: v0.6.5...v0.6.6