What is the equivalent of `nn.map_variables` in the `nnx` API? #4464

Answered by cgarciae

carlesoctav asked this question in General

carlesoctav
Dec 24, 2024

What is the equivalent of nn.map_variables in the nnx API?

Answered by cgarciae

@carlesoctav you mimic map_variables using nnx.split to get the pytree state and nnx.update to update the Module after you've modified the state. See https://flax.readthedocs.io/en/latest/nnx_basics.html#the-flax-functional-api.

View full answer

Replies: 1 comment

cgarciae
Dec 24, 2024
Maintainer

@carlesoctav you mimic map_variables using nnx.split to get the pytree state and nnx.update to update the Module after you've modified the state. See https://flax.readthedocs.io/en/latest/nnx_basics.html#the-flax-functional-api.

0 replies

Answer selected by carlesoctav

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment