How to update dictionaries in Flax when getting "ValueError: FrozenDict is immutable."?

[marcvanzee]

Original question from an internal user:

"Hello all, just checking if there is any best practice with passing and then updating dictionaries in a Flax module? I currently face a ValueError: FrozenDict is immutable."

[marcvanzee]

Flax uses FrozenDict as the datatype for dictionary, which are immutable. You should first unfreeze them using flax.core.frozen_dict.unfreeze(d). So something like this:

from flax.core.frozen_dict import unfreeze

variables = module.init(random.PRNGKey(0), ...)
# variables is now of type FrozenDict, first unfreeze it to modify it.
variables = unfreeze(variables)
# now they can be modified
variables[key] = ...