chore(deps): update dependency flax to v0.10.1 #587

renovate · 2024-07-13T23:35:10Z

This PR contains the following updates:

Package	Change	Age	Adoption	Passing	Confidence
flax	`0.8.1` -> `0.10.1`

Release Notes

google/flax (flax)

`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

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 @cgarciae in https://github.com/google/flax/pull/3758
[nnx] fix demo notebook by @cgarciae in https://github.com/google/flax/pull/3744
added nnx api reference by @chiamp in https://github.com/google/flax/pull/3762
updated rng docstring for init, apply and make_rng by @chiamp in https://github.com/google/flax/pull/3765
use note box in make_rng docstring by @cgarciae in https://github.com/google/flax/pull/3767
[nnx] improved graph update mechanism by @cgarciae in https://github.com/google/flax/pull/3759
use note box in docstrings by @chiamp in https://github.com/google/flax/pull/3769
Add reset_gate flag to MGUCell. by @carlosgmartin in https://github.com/google/flax/pull/3760
Access thread_resources via jax.interpreters.pxla instead of jax.experimental.maps by @copybara-service in https://github.com/google/flax/pull/3775
Minor doc improvements by @canyon289 in https://github.com/google/flax/pull/3588
added MGU reset_gate test by @chiamp in https://github.com/google/flax/pull/3773
[nnx] Pytrees are Trees by @cgarciae in https://github.com/google/flax/pull/3768
Use short-circuiting access to debug_key_reuse by @copybara-service in https://github.com/google/flax/pull/3781
fix tabulate on norm wrappers by @chiamp in https://github.com/google/flax/pull/3772
Add kw_only struct.dataclass test by @chiamp in https://github.com/google/flax/pull/3651
extended PyTreeNode to take dataclass kwargs by @chiamp in https://github.com/google/flax/pull/3785
[nnx] Arrays are state by @cgarciae in https://github.com/google/flax/pull/3791
[nnx] add GraphNode base class by @cgarciae in https://github.com/google/flax/pull/3790
[nnx] jit accepts many Modules by @cgarciae in https://github.com/google/flax/pull/3783
Exposing the experimental _split_transpose JAX scan parameter in Flax. by @copybara-service in https://github.com/google/flax/pull/3795
Expose nnx.GraphNode by @chiamp in https://github.com/google/flax/pull/3796
[nnx] Rngs and RngStream inherit from GraphNode by @cgarciae in https://github.com/google/flax/pull/3793
[nnx] TrainState uses struct by @cgarciae in https://github.com/google/flax/pull/3788
[nnx]

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.

renovate bot requested a review from a team as a code owner July 13, 2024 23:35

renovate bot added the dependencies Pull requests that update a dependency file label Jul 13, 2024

renovate bot force-pushed the renovate/flax-0.x-lockfile branch from cafeda6 to e803655 Compare August 16, 2024 13:51

renovate bot force-pushed the renovate/flax-0.x-lockfile branch from e803655 to 1222748 Compare August 27, 2024 19:56

renovate bot changed the title ~~chore(deps): update dependency flax to v0.8.5~~ chore(deps): update dependency flax to v0.9.0 Aug 27, 2024

renovate bot force-pushed the renovate/flax-0.x-lockfile branch from 1222748 to b2e5f65 Compare October 17, 2024 00:23

renovate bot changed the title ~~chore(deps): update dependency flax to v0.9.0~~ chore(deps): update dependency flax to v0.10.0 Oct 17, 2024

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

a4243cb

renovate bot force-pushed the renovate/flax-0.x-lockfile branch from b2e5f65 to a4243cb Compare November 1, 2024 00:10

renovate bot changed the title ~~chore(deps): update dependency flax to v0.10.0~~ chore(deps): update dependency flax to v0.10.1 Nov 1, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chore(deps): update dependency flax to v0.10.1 #587

chore(deps): update dependency flax to v0.10.1 #587

renovate bot commented Jul 13, 2024 •

edited

Loading

chore(deps): update dependency flax to v0.10.1 #587

Are you sure you want to change the base?

chore(deps): update dependency flax to v0.10.1 #587

Conversation

renovate bot commented Jul 13, 2024 • edited Loading

Release Notes

v0.10.1: Version 0.10.1

What's Changed

New Contributors

v0.10.0: Version 0.10.0

What's Changed

New Contributors

v0.9.0

What's Changed

New Contributors

v0.8.5

What's Changed

New Contributors

v0.8.4

What's Changed

v0.8.3

What's Changed

Configuration

renovate bot commented Jul 13, 2024 •

edited

Loading

`v0.10.1`: Version 0.10.1

`v0.10.0`: Version 0.10.0

`v0.9.0`

`v0.8.5`

`v0.8.4`

`v0.8.3`