How is state handled differently than parameters?

[nalzok]

After reading Stateful Computations in JAX, I have the impression that states like batch norm statistics are handled in the same fashion as model parameters: both are essentially extra input parameters to the function. If that is the case, why do we need a mutable parameter in flax.linen.Module.apply, instead of considering states as an ordinary subset part of the parameters?

In other words, what does the mutable parameter do? Despite the name, I know it does not mutate anything, since JAX arrays are immutable.

[marcvanzee]

I have the impression that states like batch norm statistics are handled in the same fashion as model parameters

No, this is not how it works. A variable dictionary in Flax contains "variable collections" as keys, which are mapped to the variables used in this collection.

• One variable collection is params, where we store the trainable parameters.
• Another is batch_stats, where we store batch statistics if using batch normalization. (If we do not use batch normalization, this collection is not present.)

The mutable argument to apply allows the user to control which variable collections are allowed to be modified during the forward pass. For instance, when you apply a module, params is not mutable because you don't modify parameters during the forward pass (you do this only in the backward pass). For batch_stats, this is different! You actually do want to modify them during the forward pass (to keep track of the mean and variance of your batch), so you need to mark this explicitly.

Despite the name, I know it does not mutate anything, since JAX arrays are immutable.

JAX arrays are immutable, so when you modify them you make a copy and use that. This is what we do in Flax as well: when you have mutable variables collections, apply actually returns a pair (model_output, modified_variable_collections):

vars_in = {'params': params, 'batch_stats': old_batch_stats}
y, mutated_vars = BN.apply(vars_in, x, mutable=['batch_stats'])
new_batch_stats = mutated_vars['batch_stats']