chore(deps): update dependency flax to v0.10.2

[renovate]

This PR contains the following updates:

Package	Change	Age	Adoption	Passing	Confidence
flax	0.8.1 -> 0.10.2

---

Release Notes

google/flax (flax)

v0.10.2: Version 0.10.2

Compare Source

What's Changed

• Add nnx.fori_loop by @​IvyZX in https://github.com/google/flax/pull/4353
• Linesearch (and lbfgs) support by @​jlperla in https://github.com/google/flax/pull/4351
• Upgrade Flax NNX Haiku Linen migration doc by @​8bitmp3 in https://github.com/google/flax/pull/4200
• Fix PRNG handling in nn.jit under nn.scan. by @​copybara-service in https://github.com/google/flax/pull/4359
• support passing arguments directly to the struct.dataclass decorator by @​copybara-service in https://github.com/google/flax/pull/4275
• Avoid assert_array_equal for PRNG keys. by @​copybara-service in https://github.com/google/flax/pull/4363
• [nnx] support pure dicts by @​cgarciae in https://github.com/google/flax/pull/4352
• [nnx] add data parallel toy example by @​cgarciae in https://github.com/google/flax/pull/4354
• Add logical axis global context support for NNX by @​IvyZX in https://github.com/google/flax/pull/4350
• [nnx] fix ToLinen kwargs by @​copybara-service in https://github.com/google/flax/pull/4270
• [nnx] use HashableMapping instead of FrozenDict by @​cgarciae in https://github.com/google/flax/pull/4376
• [nnx] fix while_loop/fori_loop bug when sharing references by @​cgarciae in https://github.com/google/flax/pull/4379
• Add flax.nnx.eval_shape docstring by @​8bitmp3 in https://github.com/google/flax/pull/4374
• Setup the flaxlib in C++, using Meson and Nanobind. by @​copybara-service in https://github.com/google/flax/pull/4380
• Add flax.nnx.remat docstring by @​8bitmp3 in https://github.com/google/flax/pull/4373
• [nnx] add checkify by @​cgarciae in https://github.com/google/flax/pull/4381
• Lint flax.nnx.while_loop docstring by @​8bitmp3 in https://github.com/google/flax/pull/4371
• Lint flax.nnx.fori_loop docstring by @​8bitmp3 in https://github.com/google/flax/pull/4370
• [nnx] add some optimizations to graph.py by @​cgarciae in https://github.com/google/flax/pull/4377
• update version to 0.10.2 by @​cgarciae in https://github.com/google/flax/pull/4387

New Contributors

• @​jlperla made their first contribution in https://github.com/google/flax/pull/4351

Full Changelog: google/flax@v0.10.1...v0.10.2

v0.10.1: Version 0.10.1

Compare Source

What's Changed

• Add Flax NNX GraphDef docstring by @​8bitmp3 in https://github.com/google/flax/pull/4302
• Flesh out the Haiku/Flax guide by @​IvyZX in https://github.com/google/flax/pull/4305
• [nnx] improve mnist tutorial by @​cgarciae in https://github.com/google/flax/pull/4316
• Update Flax Evolution from Linen to NNX guide by @​8bitmp3 in https://github.com/google/flax/pull/4289
• [nnx] try casting integers keys in State.replace_by_pure_dict by @​cgarciae in https://github.com/google/flax/pull/4317
• Fixed nnx examples bad links in the README.md by @​vfdev-5 in https://github.com/google/flax/pull/4282
• Fix philosophy link by @​jorisSchaller in https://github.com/google/flax/pull/4313
• [nnx] add gemma notebook by @​cgarciae in https://github.com/google/flax/pull/4075
• [nnx] improve init_cache docs by @​cgarciae in https://github.com/google/flax/pull/4291
• remove markdown from section titles by @​cgarciae in https://github.com/google/flax/pull/4322
• Avoid depending on JAX internals, which are about to change. by @​copybara-service in https://github.com/google/flax/pull/4326
• Remove outdated compatibility code. by @​jakevdp in https://github.com/google/flax/pull/4324
• fix ruff complaints by @​levskaya in https://github.com/google/flax/pull/4331
• Remove GeGLU activation function and golden tests. by @​copybara-service in https://github.com/google/flax/pull/4303
• Avoid using float32 in normalization for mean/var and scale/bias parameters when force_float32_reductions=False by @​copybara-service in https://github.com/google/flax/pull/4314
• Avoid assert_array_equal on PRNG keys. by @​jakevdp in https://github.com/google/flax/pull/4332
• Fix typos in Flax NNX Migrating from Haiku to Flax by @​8bitmp3 in https://github.com/google/flax/pull/4337
• Add API reference for flax.nnx.nn and improve landing page by @​IvyZX in https://github.com/google/flax/pull/4338
• [nnx] improve transforms guide by @​cgarciae in https://github.com/google/flax/pull/4333
• [nnx] cleanup gemma notebook by @​cgarciae in https://github.com/google/flax/pull/4334
• Remove non-lazy RNG compat mode and flag from flax. by @​copybara-service in https://github.com/google/flax/pull/4339
• [nnx] fix custom_vjp by @​cgarciae in https://github.com/google/flax/pull/4306
• Define model surgery in docs by @​8bitmp3 in https://github.com/google/flax/pull/4349
• [nnx] update State and variables docstrings by @​cgarciae in https://github.com/google/flax/pull/4346
• Add NNX transforms nnx.while_loop and nnx.switch by @​IvyZX in https://github.com/google/flax/pull/4343
• update version to v0.10.1 by @​cgarciae in https://github.com/google/flax/pull/4345

New Contributors

• @​vfdev-5 made their first contribution in https://github.com/google/flax/pull/4282

Full Changelog: google/flax@v0.10.0...v0.10.1

v0.10.0: Version 0.10.0

Compare Source

What's Changed

• [nnx] clear nnx basics pip logs by @​cgarciae in https://github.com/google/flax/pull/4149
• Support linen <-> nnx metadata box converging in nnx.bridge by @​IvyZX in https://github.com/google/flax/pull/4145
• Add nnx bridge API reference to site by @​IvyZX in https://github.com/google/flax/pull/4158
• [nnx] use jax-style transforms API in nnx_basics by @​cgarciae in https://github.com/google/flax/pull/4155
• [nnx] improve nnx.scan in_axes/out_axes by @​cgarciae in https://github.com/google/flax/pull/4157
• Support direct quantization for FP8 matmul by @​wenscarl in https://github.com/google/flax/pull/3922
• Upgrade Flax NNX Model Surgery by @​8bitmp3 in https://github.com/google/flax/pull/4135
• [nnx] add more Variable proxy methods by @​cgarciae in https://github.com/google/flax/pull/4170
• [nnx] disallow Array leaves by @​copybara-service in https://github.com/google/flax/pull/4172
• Internal change by @​copybara-service in https://github.com/google/flax/pull/4176
• [nnx] improve landing page and nnx_basics messaging by @​cgarciae in https://github.com/google/flax/pull/4168
• Fixes a small bug in flax.linen.share_scope, where the scopes of children of the module being merged that were created before setup(),were not being updated to point to the new scope, and so they would end up staying under the original tree. by @​copybara-service in https://github.com/google/flax/pull/4150
• Move all NNX content up a level to be equal with Linen, to make python packaging more consistent. by @​copybara-service in https://github.com/google/flax/pull/4177
• Add a guide for nnx.bridge by @​IvyZX in https://github.com/google/flax/pull/4171
• [nnx] improve Optimizer metadata propagation by @​cgarciae in https://github.com/google/flax/pull/4180
• [nnx] enable sharding transformation on integer prefixes by @​cgarciae in https://github.com/google/flax/pull/4185
• Support linen.LogicallyPartitioned <-> nnx.Variable by @​IvyZX in https://github.com/google/flax/pull/4161
• Clean up axis hooks in nnx.Variable by @​IvyZX in https://github.com/google/flax/pull/4189
• Merge nnx.errors to flax.errors by @​IvyZX in https://github.com/google/flax/pull/4186
• [nnx] optimize jit by @​cgarciae in https://github.com/google/flax/pull/4191
• Split documentation for Linen and NNX by @​cgarciae in https://github.com/google/flax/pull/4192
• Partially revert #​4192 which sets back a bunch of previous merged pushes. by @​copybara-service in https://github.com/google/flax/pull/4201
• Align bridge variable tree structures by @​IvyZX in https://github.com/google/flax/pull/4194
• [NNX site] Fix landing page and banner phrasing and add examples page by @​IvyZX in https://github.com/google/flax/pull/4202
• shorten banners by @​cgarciae in https://github.com/google/flax/pull/4206
• Add trimmed Linen to NNX guide by @​IvyZX in https://github.com/google/flax/pull/4209
• Minor documentation fixes for AxisMetadata. by @​copybara-service in https://github.com/google/flax/pull/4178
• fix tests for numpy 2.0 compatibility by @​copybara-service in https://github.com/google/flax/pull/4215
• Forward all arguments when using nnx.transforms.deprecated.scan as a decorator. by @​copybara-service in https://github.com/google/flax/pull/4208
• [nnx] add transforms guide by @​cgarciae in https://github.com/google/flax/pull/4197
• [nnx] fix transforms guide by @​cgarciae in https://github.com/google/flax/pull/4223
• Flax NNX GSPMD guide by @​IvyZX in https://github.com/google/flax/pull/4220
• Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. by @​copybara-service in https://github.com/google/flax/pull/4225
• [nnx] add Randomness guide by @​cgarciae in https://github.com/google/flax/pull/4216
• Add pure dict conversion util functions to nnx.State. by @​IvyZX in https://github.com/google/flax/pull/4230
• [nnx] Simplify traversal by @​cgarciae in https://github.com/google/flax/pull/4205
• Fix false positive tracer leaks in flax library. by @​copybara-service in https://github.com/google/flax/pull/4232
• [nnx] add flaxlib by @​copybara-service in https://github.com/google/flax/pull/4235
• [nnx] improve docs by @​cgarciae in https://github.com/google/flax/pull/4236
• point nnx banner to flax-linen by @​cgarciae in https://github.com/google/flax/pull/4237
• update banners by @​cgarciae in https://github.com/google/flax/pull/4238
• Fix scale dtype and refactor q_dot_dq by @​wenscarl in https://github.com/google/flax/pull/4229
• update banners by @​cgarciae in https://github.com/google/flax/pull/4241
• Add redirects for Linen guide links in the NNX site scope. by @​IvyZX in https://github.com/google/flax/pull/4242
• Internal change by @​copybara-service in https://github.com/google/flax/pull/4243
• Copybara import of the project: by @​copybara-service in https://github.com/google/flax/pull/4245
• Update Flax NNX Scale Up SPMD guide by @​8bitmp3 in https://github.com/google/flax/pull/4239
• Upgrade Flax NNX basics doc by @​8bitmp3 in https://github.com/google/flax/pull/4173
• Improve landing page, glossary and misc by @​IvyZX in https://github.com/google/flax/pull/4244
• Nitting and adding links by @​8bitmp3 in https://github.com/google/flax/pull/4248
• enable doctest on notebooks by @​cgarciae in https://github.com/google/flax/pull/4250
• Update index.rst by @​ariG23498 in https://github.com/google/flax/pull/4251
• Add NNX checkpointing guide by @​IvyZX in https://github.com/google/flax/pull/4249
• Add checkpointing guide to website index. by @​copybara-service in https://github.com/google/flax/pull/4263
• Update to Flax NNX Transforms doc by @​8bitmp3 in https://github.com/google/flax/pull/4264
• Add why nnx by @​cgarciae in https://github.com/google/flax/pull/4240
• [nnx] add cloudpickle support by @​cgarciae in https://github.com/google/flax/pull/4253
• Fix typo: impost to import by @​Vilin97 in https://github.com/google/flax/pull/4256
• [nnx] revive TrainState toy example by @​cgarciae in https://github.com/google/flax/pull/4226
• [nnx] add custom_vjp to docs by @​cgarciae in https://github.com/google/flax/pull/4266
• remove flax-nnx urls by @​cgarciae in https://github.com/google/flax/pull/4267
• Add flatten to nnx.graph autosummary in graph.rst by @​8bitmp3 in https://github.com/google/flax/pull/4255
• [nnx] add FSDP toy example with custom optimizer by @​cgarciae in https://github.com/google/flax/pull/4183
• Update Flax NNX Landing Page by @​8bitmp3 in https://github.com/google/flax/pull/4274
• Update to Flax NNX Model Surgery by @​8bitmp3 in https://github.com/google/flax/pull/4276
• Update Why Flax NNX guide by @​8bitmp3 in https://github.com/google/flax/pull/4262
• Update to Flax NNX MNIST tutorial by @​8bitmp3 in https://github.com/google/flax/pull/4277
• [nnx] improve randomness guide by @​cgarciae in https://github.com/google/flax/pull/4281
• Remove notebook exceptions in docs_nnx doctest by @​IvyZX in https://github.com/google/flax/pull/4285
• [nnx] add PrefixMapping by @​cgarciae in https://github.com/google/flax/pull/4278
• [nnx] state filters by @​cgarciae in https://github.com/google/flax/pull/4288
• Fix devcontainer setup by @​jorisSchaller in https://github.com/google/flax/pull/4299
• Ugrade Flax NNX Checkpointing guide by @​8bitmp3 in https://github.com/google/flax/pull/4294
• Update Flax NNX Scale Up guide by @​8bitmp3 in https://github.com/google/flax/pull/4296
• Porting RNN from Linen to NNX by @​zinccat in https://github.com/google/flax/pull/4272
• Update Flax NNX Glossary by @​8bitmp3 in https://github.com/google/flax/pull/4284
• update version to 0.10.0 by @​cgarciae in https://github.com/google/flax/pull/4292

New Contributors

• @​ariG23498 made their first contribution in https://github.com/google/flax/pull/4251
• @​Vilin97 made their first contribution in https://github.com/google/flax/pull/4256
• @​jorisSchaller made their first contribution in https://github.com/google/flax/pull/4299
• @​zinccat made their first contribution in https://github.com/google/flax/pull/4272

Full Changelog: google/flax@v0.9.0...v0.10.0

v0.9.0

Compare Source

What's Changed

• Add NNX surgery guide by @​IvyZX in https://github.com/google/flax/pull/4005
• Port gemma/transformer to NNX by @​copybara-service in https://github.com/google/flax/pull/4019
• upgrade python to 3.10 + use pyupgrade by @​cgarciae in https://github.com/google/flax/pull/4038
• [nnx] add Using Filters guide by @​cgarciae in https://github.com/google/flax/pull/4028
• v0.8.6 by @​cgarciae in https://github.com/google/flax/pull/4040
• allow imagenet training profiling to be disabled in config by @​copybara-service in https://github.com/google/flax/pull/4043
• [nnx] LoRAParam inherits from Param by @​cgarciae in https://github.com/google/flax/pull/3988
• [linen] allows multiple compact methods by @​cgarciae in https://github.com/google/flax/pull/3808
• Added support of NANOO fp8. by @​wenchenvincent in https://github.com/google/flax/pull/3993
• Add functool.wraps() annotation to flax.nn.jit. by @​copybara-service in https://github.com/google/flax/pull/4051
• Fix typo in nnx_basics doc by @​rajasekharporeddy in https://github.com/google/flax/pull/4047
• [nnx] fix Variable overloads and add shape/dtype properties by @​cgarciae in https://github.com/google/flax/pull/4049
• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @​copybara-service in https://github.com/google/flax/pull/4039
• [nnx] stabilize unsafe_pytree by @​cgarciae in https://github.com/google/flax/pull/4030
• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @​copybara-service in https://github.com/google/flax/pull/4055
• [NVIDIA] Rename fp8 custom dtype to fp32_max_grad by @​kaixih in https://github.com/google/flax/pull/3984
• [nnx] fix mnist_tutorial colab link by @​cgarciae in https://github.com/google/flax/pull/4063
• [nnx] fix Accuracy on eager mode by @​cgarciae in https://github.com/google/flax/pull/4065
• Update orbax_upgrade_guide.rst for async checkpointing usage examples by @​kaushaladiti-2802 in https://github.com/google/flax/pull/4036
• Re-enable some tests after Python 3.9 is dropped by @​IvyZX in https://github.com/google/flax/pull/4067
• Rename nnx.compat to nnx.bridge by @​IvyZX in https://github.com/google/flax/pull/4066
• [nnx] improve mnist tutorial by @​cgarciae in https://github.com/google/flax/pull/4070
• Modify Flax checkpointing in preparation for cl/650338576. by @​copybara-service in https://github.com/google/flax/pull/4072
• Remove some outdated backward-compatibility code. by @​copybara-service in https://github.com/google/flax/pull/4068
• [NVIDIA] Add a user guide for fp8 by @​kaixih in https://github.com/google/flax/pull/4076
• [nnx] add extract APIs by @​cgarciae in https://github.com/google/flax/pull/4078
• [example]: remove lm1b useless parallism rules by @​knightXun in https://github.com/google/flax/pull/4077
• [nnx] improve filters guide by @​cgarciae in https://github.com/google/flax/pull/4059
• [nnx] add call by @​cgarciae in https://github.com/google/flax/pull/4004
• Ignore Orbax warning in deprecated flax.training.checkpoints.py to unbreak head doctest by @​IvyZX in https://github.com/google/flax/pull/4092
• fix mypy failures due tu numpy update by @​cgarciae in https://github.com/google/flax/pull/4098
• [linen] generalize transform caching by @​copybara-service in https://github.com/google/flax/pull/4057
• [linen] fold rngs on jit to improve caching by @​copybara-service in https://github.com/google/flax/pull/4064
• Add shape-based lazy init to LinenToNNX (prev LinenWrapper) by @​IvyZX in https://github.com/google/flax/pull/4081
• [nnx] add reseed by @​cgarciae in https://github.com/google/flax/pull/4099
• [nnx] add split/merge_inputs by @​cgarciae in https://github.com/google/flax/pull/4084
• Perform shape checks for self.param AFTER unboxing by @​danielwatson6 in https://github.com/google/flax/pull/4079
• fix restore_checkpoint example in docstring by @​copybara-service in https://github.com/google/flax/pull/4101
• [numpy] Fix users of NumPy APIs that are removed in NumPy 2.0. by @​copybara-service in https://github.com/google/flax/pull/4104
• set profile_duration_ms = None as in periodic_actions there's default value for both num_profile_steps and profile_duration_ms, and the profile stopping condition is when both num_profile_steps and profile_duration_ms are satisfied, so setting profile_duration_ms=None so that the passed num_profile_steps value gets used by @​copybara-service in https://github.com/google/flax/pull/4096
• [linen] add share_scope by @​cgarciae in https://github.com/google/flax/pull/4102
• Allow metadata pass-through in flax.struct.field by @​cool-RR in https://github.com/google/flax/pull/4056
• avoid mixing einsum_dot_general and einsum argument by specifying them explicitly in the caller. by @​copybara-service in https://github.com/google/flax/pull/4115
• Add logging to track deprecated codepaths. by @​copybara-service in https://github.com/google/flax/pull/4121
• [pmap no rank reduce cleanup]: When flipping the by @​copybara-service in https://github.com/google/flax/pull/4125
• Add NNXToLinen wrapper to nnx.bridge by @​IvyZX in https://github.com/google/flax/pull/4126
• Switch NNX to use Treescope instead of Penzai. by @​copybara-service in https://github.com/google/flax/pull/4132
• Add GroupNorm to NNX normalization layers by @​treigerm in https://github.com/google/flax/pull/4095
• [nnx] fix initializing propagation by @​cgarciae in https://github.com/google/flax/pull/4134
• add JAX-style NNX Transforms FLIP by @​cgarciae in https://github.com/google/flax/pull/4108
• Fix _ParentType annotation by @​dcharatan in https://github.com/google/flax/pull/4120
• add uv.lock file by @​copybara-service in https://github.com/google/flax/pull/4139
• use uv package manager by @​cgarciae in https://github.com/google/flax/pull/4136
• More testing and misc fixes on wrappers by @​IvyZX in https://github.com/google/flax/pull/4137
• Fix link to orbax documentation by @​cool-RR in https://github.com/google/flax/pull/4123
• [nnx] experimental transforms by @​cgarciae in https://github.com/google/flax/pull/3963
• [nnx] improve docs by @​cgarciae in https://github.com/google/flax/pull/4141
• remove repeated license headers by @​cgarciae in https://github.com/google/flax/pull/4148
• update Flax to version 0.9.0 by @​copybara-service in https://github.com/google/flax/pull/4147

New Contributors

• @​wenchenvincent made their first contribution in https://github.com/google/flax/pull/3993
• @​rajasekharporeddy made their first contribution in https://github.com/google/flax/pull/4047
• @​kaushaladiti-2802 made their first contribution in https://github.com/google/flax/pull/4036
• @​knightXun made their first contribution in https://github.com/google/flax/pull/4077
• @​danielwatson6 made their first contribution in https://github.com/google/flax/pull/4079
• @​cool-RR made their first contribution in https://github.com/google/flax/pull/4056
• @​treigerm made their first contribution in https://github.com/google/flax/pull/4095
• @​dcharatan made their first contribution in https://github.com/google/flax/pull/4120

Full Changelog: google/flax@v0.8.5...v0.9.0

v0.8.5

Compare Source

What's Changed

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

New Contributors

• @​mcsmart76 made their first contribution in https://github.com/google/flax/pull/3953
• @​jkarwowski made their first contribution in https://github.com/google/flax/pull/3982
• @​sanderland made their first contribution in https://github.com/google/flax/pull/4002

Full Changelog: google/flax@v0.8.4...v0.8.5

v0.8.4

Compare Source

What's Changed

• fixed codecov by @​chiamp in https://github.com/google/flax/pull/3895
• Make FlatState a Mapping instead of a dict by @​NeilGirdhar in https://github.com/google/flax/pull/3880
• Share nnx node registry between threads by @​NeilGirdhar in https://github.com/google/flax/pull/3901
• fixed jnp.clip deprecation by @​chiamp in https://github.com/google/flax/pull/3905
• Added three tab option to sphinx directive codediff and added testing for first tab by @​chiamp in https://github.com/google/flax/pull/3847
• Add support for jax.sharding.PartitionSpec.UNCONSTRAINED in logical specification by @​copybara-service in https://github.com/google/flax/pull/3902
• [nnx] fix mypy and pytype by @​cgarciae in https://github.com/google/flax/pull/3894
• [nnx] fix iter_nodes by @​cgarciae in https://github.com/google/flax/pull/3889
• [nnx] Sequential uses regular list by @​cgarciae in https://github.com/google/flax/pull/3909
• [nnx] add ConvTranspose by @​cgarciae in https://github.com/google/flax/pull/3908
• [nnx] add Module pytree_experimental static test by @​cgarciae in https://github.com/google/flax/pull/3864
• Added docstring for Module.scope.path by @​chiamp in https://github.com/google/flax/pull/3913
• [linen] test jit caching with state updates by @​cgarciae in https://github.com/google/flax/pull/3900
• v0.8.4 by @​cgarciae in https://github.com/google/flax/pull/3891
• [linen] enable separate initializers for out layer in MultiHeadDotProductAttention by @​cgarciae in https://github.com/google/flax/pull/3835
• [nnx] cleanup graph by @​cgarciae in https://github.com/google/flax/pull/3915
• [nnx] fix bugs by @​cgarciae in https://github.com/google/flax/pull/3925
• Replace deprecated jax.tree_* functions with jax.tree.* by @​copybara-service in https://github.com/google/flax/pull/3926
• [nnx] Object refactor by @​cgarciae in https://github.com/google/flax/pull/3910
• [nnx] add iter_graph by @​cgarciae in https://github.com/google/flax/pull/3919
• [nnx] add compat by @​cgarciae in https://github.com/google/flax/pull/3921
• [nnx] transforms refactor by @​cgarciae in https://github.com/google/flax/pull/3927
• added equivalence test for nnx.ConvTranspose by @​chiamp in https://github.com/google/flax/pull/3934
• added equivalence test for nnx.Sequential by @​chiamp in https://github.com/google/flax/pull/3935
• [NNX] Add LoRA and LoRALinear to NNX by @​IvyZX in https://github.com/google/flax/pull/3929
• [nnx] fix substate mutability by @​cgarciae in https://github.com/google/flax/pull/3932
• [nnx] improve update context by @​cgarciae in https://github.com/google/flax/pull/3933
• [nnx] move out of experimental by @​cgarciae in https://github.com/google/flax/pull/3936

Full Changelog: google/flax@v0.8.3...v0.8.4

v0.8.3

Compare Source

What's Changed

• Add git fetch upstream to contributing doc. by @​carlosgmartin in https://github.com/google/flax/pull/3757
• removed getattr/setattr unboxing magic from nnx.Pytree by @​chiamp in https://github.com/google/flax/pull/3743
• added Einsum layer to NNX by @​chiamp in https://github.com/google/flax/pull/3741
• Make TrainState's step possibly jax.Array. This makes replicate valid for type checking. by @​copybara-service in https://github.com/google/flax/pull/3763
• v0.8.3 by [@​cga

---

Configuration

📅 Schedule: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 Automerge: Disabled by config. Please merge this manually once you are satisfied.

♻ Rebasing: Whenever PR becomes conflicted, or you tick the rebase/retry checkbox.

🔕 Ignore: Close this PR and you won't be reminded about this update again.

---

• If you want to rebase/retry this PR, check this box

---

This PR was generated by Mend Renovate. View the repository job log.