google/flax v0.10.7

0.10.7

on GitHub

What's Changed

• Added identity export from JAX by @jlperla in #4652
• Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning. by @copybara-service in #4727
• fix merge by @copybara-service in #4731
• [nnx] make Variable a pytree by @cgarciae in #4728
• [nnx] add JitWrapped API by @cgarciae in #4699
• Update JAX nightly index usage by @copybara-service in #4733
• [nnx] mutable array p1 by @cgarciae in #4715
• add dataclass by @copybara-service in #4739
• [flax] unconditionally register nnx.Variable as a pytree by @copybara-service in #4748
• Updated version of pre-commit-hooks in .pre-commit-config.yaml by @vfdev-5 in #4746
• Fixed docstring visibility for nnx.eval_shape by @vfdev-5 in #4747
• Added keep_rngs arg to MHA to optionally store rngs by @vfdev-5 in #4749
• MultiHeadAttention only keeps rngs if dropout_rate is positive by @copybara-service in #4750
• [nnx] mutable array p2 by @cgarciae in #4741
• Add in_kv_features argument to nnx.MultiHeadAttention, addressing #4756. by @copybara-service in #4757
• Fix broken link for Transforms guide by @nireekshak in #4763
• Minor improvements of lm1b_nnx example by @vfdev-5 in #4745
• Fix head CI tests by @IvyZX in #4764
• Fix typos by @nireekshak in #4725
• Check for leaves of type variablelib.Variable when getting sharding specs. by @copybara-service in #4769
• Fixes #1925 non-str dict keys not suppoted in module state by @muhrin in #4563
• Modified the Functional API link by @nireekshak in #4767
• Fix hardcoded link to filter guide in docs by @hamogu in #4768
• Fix bad doc links by @IvyZX in #4770
• revise axes_scan to flatten argument pytrees only once by @copybara-service in #4772
• Simplify ToNNX access of Linen module methods by @IvyZX in #4766
• Use .input_formats and .output_formats in place of .input_layouts and .output_layouts respectively. by @copybara-service in #4784
• Exposed OptState in nnx module by @vfdev-5 in #4788
• Fixes colab link for nnx docs by @vfdev-5 in #4775
• Internal changes by @copybara-service in #4786
• Fix typo in Flax nnx_basics doc. by @copybara-service in #4781
• update version to 0.10.7 by @cgarciae in #4798

New Contributors

• @nireekshak made their first contribution in #4763
• @muhrin made their first contribution in #4563
• @hamogu made their first contribution in #4768

Full Changelog: v0.10.6...v0.10.7