Align bridge variable tree structures #4194
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
IvyZX
force-pushed
the
bdg-tree
branch
3 times, most recently
from
8eceb3c to
172e832
Compare
cgarciae
reviewed
Sep 14, 2024
flax/nnx/bridge/wrappers.py
Outdated
| import dataclasses | ||
| import typing as tp | ||
| from typing import Any | ||
|
|
||
| from flax import nnx | ||
| from flax import linen | ||
| from flax import traverse_util |
There was a problem hiding this comment.
NIT: flax.nnx.traversals has similar same APIs but with more accurate typing.
cgarciae
reviewed
Sep 14, 2024
flax/nnx/bridge/wrappers.py
Outdated
| @@ -85,6 +87,50 @@ def lazy_init(fn: Module | tp.Callable[..., tp.Any], *args, **kwargs): | |||
| return fn | |||
|
|
|||
|
|
|||
| def _recursive_merge(dict1, dict2): | |||
There was a problem hiding this comment.
I think this could be expressed using traversals / traverse_util as:
flat_map = traversal.flatten_mapping(dict1)
flat_map |= traversal.flatten_mapping(dict2)
return traversal.unflatten_mapping(flat_map)Not sure if there are edge cases though but it seems easy.
There was a problem hiding this comment.
Looks good! Adopted.
cgarciae
reviewed
Sep 14, 2024
flax/nnx/bridge/wrappers.py
Outdated
| dict1[key] = _recursive_merge(dict1[key], value) | ||
| else: | ||
| # Merge non-dictionary values | ||
| dict1[key] = value |
There was a problem hiding this comment.
Is mutating dict1 always safe?
There was a problem hiding this comment.
I think using nnx.traversals is indeed easier!
cgarciae
reviewed
Sep 14, 2024
flax/nnx/bridge/wrappers.py
Outdated
| def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict: | ||
| linen_structured = {} | ||
| for kp, v in traverse_util.flatten_dict( | ||
| nnx_attrs, is_leaf=lambda _, x: isinstance(x, nnx.Variable | nnx.GraphDef)).items(): |
There was a problem hiding this comment.
I didn't know Unions were now valid in isinstance, nice!
cgarciae
approved these changes
Sep 16, 2024
IvyZX
force-pushed
the
bdg-tree
branch
2 times, most recently
from
55cbdef to
cac253e
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Fixing bridge API in a bunch of ways:
ToNNXwill convert the whole variables structure to NNX style. If your underlying Linen module has variablefooat collectionbar, itsToNNXversion will have an attributefoowith typebar, instead of an attributebarwith a dict{'foo': ...}.This means you can freely put
ToNNXin the top or middle or back of the whole model layer stack, and the weight pytree structure shouldn't change.Same goes for
ToLinen- if your top-level type is Linen, the whole variable tree shall be Linen-style.If you have a vanilla
nnx.Variablewith no sharding metadata, hooks, etc,ToLinenwill not convert it into anNNXMeta, but instead just keep the vanilla JAX array inside. This makes it more intuitive and pytree-structure-proof for any Linen users not using partitioning metadata.nn.get_partition_specnow works onNNXMetawrappers, and any other wrapper that hasget_partition_specmethod.Updated the
nnx.bridgeguide accordingly.