chore(deps): update dependency flax to v0.11.2 #587
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Coming soon: The Renovate bot (GitHub App) will be renamed to Mend. PRs from Renovate will soon appear from 'Mend'. Learn more here.
This PR contains the following updates:
0.8.1
->0.11.2
Warning
Some dependencies could not be looked up. Check the Dependency Dashboard for more information.
Release Notes
google/flax (flax)
v0.11.2
: 0.11.2Compare Source
What's Changed
nnx.merge
now doesn't create a copy of the Variables in the incoming states by default, meaning that the new merged structures holds references to the incoming Variables. This enables new patterns, for example its now possible to create models with the same state but with different runtime behavior:model
andeval_model
share the same Variables and are therefore kept in sync but have different runtime behavior, this avoids having to constantly mutate a single model back and forth between different runtime modes which can be error prone / cause unwanted recompilation.To keep the old behavior use
nnx.merge(..., copy=True)
.PRs
New Contributors
Full Changelog: google/flax@v0.11.1...v0.11.2
v0.11.1
Compare Source
What's Changed
Sequential()
be identity by @SobhanMP in #4796jax.sharding.use_mesh
withjax.set_mesh
.jax.set_mesh
can act as a global setter or a context manager. by @copybara-service[bot] in #4862New Contributors
Full Changelog: google/flax@v0.11.0...v0.11.1
v0.11.0
Compare Source
v0.11.0 - Pytrees, MutableArrays, and more!
This version of Flax introduces some changes to improve interop with native JAX and adds support for the new
jax.experimental.MutableArray
. More on this soon! However, some breaking changes to align with the JAX way of doing things were necessary. Most code should remain intact, however, the following changes deviate from the current behavior:Rngs
in standard layers: all standard layers no longer hold a shared reference to therngs
object given in the constructor, instead they now keep afork
-ed copy of theRngs
orRngStream
objects. This impacts Using Rngs in NNX Transforms and Loading Checkpoints with RNGs.model
to avoid reference sharing, instead themodel
must be provided as the first argument toupdate
.split
andmerge
when interacting trivially with raw JAX transforms (state must still be manually propagated if not using MutableArrays, and referential transparency is still an issue). This affects when operating on Pytrees containing NNX Objects withjax.tree.*
APIs.Checkout the full NNX 0.10 to NNX 0.11 migration guide.
In the near future we'll share more information about new ways of using NNX with JAX transforms directly by leveraging the new Pytree and MutableArray support. Stay tuned!
What's Changed
.type
usage by @vfdev-5 in #4823.value
to[...]
in modules_test.py by @lukeyeh in #4815transforms_test.py
from.value
to[...]
by @lukeyeh in #4841New Contributors
Full Changelog: google/flax@v0.10.7...v0.11.0
v0.10.7
: 0.10.7Compare Source
What's Changed
.input_formats
and.output_formats
in place of.input_layouts
and.output_layouts
respectively. by @copybara-service in #4784nnx_basics
doc. by @copybara-service in #4781New Contributors
Full Changelog: google/flax@v0.10.6...v0.10.7
v0.10.6
: 0.10.6Compare Source
What's Changed
attention_bias
parameter toMultiHeadDotProductAttention
. by @copybara-service in #4694attention_bias
parameter toMultiHeadDotProductAttention
. Add parameter to all overloads to make pytype happy. by @copybara-service in #4702Full Changelog: google/flax@v0.10.5...v0.10.6
v0.10.5
: 0.10.5Compare Source
What's Changed
wrappers_test.py
to another file. by @copybara-service in #4581beam_search
loop. by @copybara-service in #4615New Contributors
Full Changelog: google/flax@v0.10.4...v0.10.5
v0.10.4
: Release 0.10.4Compare Source
What's Changed
is_initializing
API by @copybara-service in #4550nnx.State
methods by @IvyZX in #4561New Contributors
Full Changelog: google/flax@v0.10.3...v0.10.4
v0.10.3
: Version 0.10.3Compare Source
What's Changed
async_checkpointer.py
reference by @emmanuel-ferdman in #4385Why Flax NNX
documentation by @rajasekharporeddy in #4425merge
docs in graph.py by @8bitmp3 in #4411nnx.bridge.variables.nnx_attrs_to_linen_vars
takennx.VariableState
as argument. by @copybara-service in #4473Param(None)
lines from NNX by @IvyZX in #4504nnx.Module.perturb
by @IvyZX in #4515New Contributors
Full Changelog: google/flax@v0.10.2...v0.10.3
v0.10.2
: Version 0.10.2Compare Source
What's Changed
nnx.fori_loop
by @IvyZX in #4353nn.jit
undernn.scan
. by @copybara-service in #4359flax.nnx.eval_shape
docstring by @8bitmp3 in #4374flax.nnx.remat
docstring by @8bitmp3 in #4373New Contributors
Full Changelog: google/flax@v0.10.1...v0.10.2
v0.10.1
: Version 0.10.1Compare Source
What's Changed
nnx.while_loop
andnnx.switch
by @IvyZX in #4343New Contributors
Full Changelog: google/flax@v0.10.0...v0.10.1
v0.10.0
: Version 0.10.0Compare Source
What's Changed
nnx.bridge
by @IvyZX in #4145Configuration
📅 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.
This PR was generated by Mend Renovate. View the repository job log.