v0.7.5

What's Changed

• Add method-to-model section to Haiku migration guide by @IvyZX in #3277
• updated haiku guide with new JAX RNG api by @chiamp in #3343
• changed resnet v1 to v1.5 by @chiamp in #3344
• updated flax_basics by @chiamp in #3342
• Add find methods and magic methods for Cursor API by @chiamp in #3306
• fix DeprecationWarnings by @chiamp in #3352
• Make checkpoint guide path absolute by @IvyZX in #3358
• use jax.Array type for rng keys by @chiamp in #3354
• Disable pyink by @cgarciae in #3356
• removed DeprecationWarning filter by @chiamp in #3359
• Don't propagate default args in Tabulate by @cgarciae in #3357
• added spectral norm by @chiamp in #3335
• fixed bind-unbind bug by @chiamp in #3365
• Add MaxText open source LLM to index.rst by @8bitmp3 in #3368
• fixed typo by @chiamp in #3367
• Add fp8 custom op and unit test by @wenscarl in #3322
• fixed scope typing by @chiamp in #3371
• Trailing whitespace fixes. by @levskaya in #3373
• Make fp8 ops use explicit broadcasting. by @levskaya in #3374
• Conditonal contraints for numpy and clu by @cgarciae in #3394
• add truncated_normal to initializers by @levskaya in #3401
• fix HEAD by @chiamp in #3404
• updated imagenet readme by @chiamp in #3383
• Add Flax FAQ - how to search, how to take the derivative w.r.t. a hidden layer, remat_scan vs scan(remat), recommended training libraries/metrics by @8bitmp3 in #3267
• added rmsnorm to api docs by @chiamp in #3406
• updated docs by @chiamp in #3370
• fix HEAD by @chiamp in #3408
• added weightnorm layer by @chiamp in #3405
• added attention refactor to changelog by @chiamp in #3412
• added dropout arg to MultiHeadDotProductAttention by @chiamp in #3384
• fix spectralnorm layer by @chiamp in #3403
• remove pdf target by @cgarciae in #3415
• added import precommit hook by @chiamp in #3410
• fixed GRU docstring by @chiamp in #3419
• Replaces pjit with jit in spmd.py by @copybara-service in #3421
• Ignore transient chex deprecationwarning. by @levskaya in #3427
• Remove transformers dependency from docs by @cgarciae in #3431
• updated pytorch upgrade guide by @chiamp in #3432
• Simplify abstract rng creation in param shape check. by @levskaya in #3429
• [NVIDIA] Update the FP8 support by @kaixih in #3435
• fixed mypy errors at HEAD by @chiamp in #3440
• added has_improved field to EarlyStopping by @chiamp in #3385
• fully deprecated old RNN api by @chiamp in #3425
• updated lm1b example with jit by @chiamp in #3302
• added MGU class by @chiamp in #3418
• Make Flax Basics visible by @8bitmp3 in #3443
• update for release v0.7.5 by @levskaya in #3444

New Contributors

• @wenscarl made their first contribution in #3322
• @kaixih made their first contribution in #3435

Full Changelog: v0.7.4...v0.7.5

v0.7.4

What's Changed

Added python version constraint >=3.9.

• Bump version no by 1 post-release. by @copybara-service in #3328
• Update minimum python version info. by @levskaya in #3331

Full Changelog: v0.7.3...v0.7.4

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