What is the proper way to update module variables defined in setup?

[minqi]

What is the proper way to update variables defined in setup? There is no documentation that I could find that describes this.

My use case is that I have two submodules A and B, where A and B have the same architecture. However, I want to update B using a momentum update of the weights of A, i.e. m*W_B + (1-m)*W_A, for some 0 < m < 1.

I am currently storing the parameters of module B in a variable. However, when I try to update this value inside __call__ as triggered by init, I receive the error ValueError: FrozenDict is immutable. This occurs even when I pass mutable=True when calling init.

Any pointers or suggestions on how to accomplish this would be appreciated.

[cgarciae]

Hey @minqi, can you post a snippet? Personally I've done this outside of apply e.g:

state = train_step(state, inputs, ...)
ema_params = momentum_update(state.params, ema_params)
...
preds = state.apply_fn({'params': ema_params}, inputs, ...)