google/flax v0.8.4

on GitHub

What's Changed

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

Full Changelog: v0.8.3...v0.8.4