Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[nnx] support pure dicts #4352

Merged
merged 1 commit into from
Nov 8, 2024
Merged

[nnx] support pure dicts #4352

merged 1 commit into from
Nov 8, 2024

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Oct 31, 2024 •
edited
Loading

What does this PR do?

  • nnx.merge and nnx.update can now use pure dictionaries returned from State.to_pure_dict.
  • NodeDef.leaves now contain a new type called VariableDef which contains the static information of Variables, this makes it possible to reconstruct a Variables from states that dont contain VariableStates as its the case with dictionaries returned by to_pure_dict.
  • Does some minor refactors.

Example:

m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(m)

pure_state = state.to_pure_dict() # remove leaf metdata

m2 = nnx.merge(graphdef, pure_state) # merge from pure state

@cgarciae cgarciae force-pushed the nnx-pure-dict-support branch 4 times, most recently from 3872314 to 627bcdc Compare November 1, 2024 13:11
IvyZX
IvyZX approved these changes Nov 7, 2024
View reviewed changes
Copy link
Collaborator

@IvyZX IvyZX left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Super cool!

examples/nnx_toy_examples/02_lifted_transforms.py Outdated
@@ -86,6 +86,7 @@ def test_step(model: MLP, batch):
total_steps = 10_000
for step, batch in enumerate(dataset(32)):
train_step(model, optimizer, batch)
print(nnx.graph.GRAPH_CONTEXT)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: remove print

flax/nnx/graph.py Outdated
value = value.to_variable()
children[key] = value
elif noderef.index in index_ref:
# if not is_state_leaf(value):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: remove comments

flax/nnx/graph.py Outdated
f'Expected a Variable type for {key!r}, but got {type(value)}.'
)
assert isinstance(variabledef, VariableDef)
# if not isinstance(value, VariableState):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Same here

@copybara-service copybara-service bot merged commit 0fabd3f into main Nov 8, 2024
19 checks passed
@copybara-service copybara-service bot deleted the nnx-pure-dict-support branch November 8, 2024 21:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@IvyZX IvyZX IvyZX approved these changes

Assignees
No one assigned
Labels
pull ready
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@cgarciae @IvyZX