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
What does this PR do?
nnx.mergeandnnx.updatecan now use pure dictionaries returned fromState.to_pure_dict.NodeDef.leavesnow contain a new type calledVariableDefwhich contains the static information ofVariables, this makes it possible to reconstruct aVariablesfrom states that dont containVariableStates as its the case with dictionaries returned byto_pure_dict.Example: