split params and variables in .init()

[GJBoth]

Hi all,

I'm working with models with states and parameters and having to manually split the parameters and variables, and then delete the original variables feels not very streamlined. Would it be an idea to add an option split_params=True/False to the init function, which if true would return a tuple (params, variables)?

[GJBoth]

Related to this: the recurrent examples don't initialize the state as a variable, but as a separate carry. Is there any specific reason for this?

[marcvanzee]

Could you be a bit more specific? Which example, and which initialization are you referring to?

[GJBoth]

Sorry, my bad. I'm talking about the seq2seq example: https://github.com/google/flax/blob/master/examples/seq2seq/train.py
More specifically, I'm talking about LSTMCell which has an initialize carry:

flax/flax/linen/recurrent.py

Line 66 in 6c14438

class LSTMCell(RNNCellBase):

Why not initialize it using self.variable and carrying it through an argument to nn.transforms.scan?

[marcvanzee]

I guess that would be possible as well, but I don't see why that would be useful.

We generally use variables to either optimize something (params), or for keeping track of some statistics we want to be able to use later on (e.g., batch_stats). In this specific use-case it doesn't seem useful to store the carry in the variables, since they are only used when applying scan and can be discarded afterwards.

[GJBoth]

So the choice whether to put something in a self.variable depends on if it you would use it after the scan? Makes sense... In my case (echo-state networks) I actually need the state after the scan, but i'll compare both approaches to see which one is more clear. Thanks!

[jheek]

Currently the easiest thing to split of the params is to call FrozenDict.pop:

variables = Model().init(...)
state, params = variables.pop("params")

That's pretty short and you can even just write the one-liner: Model().init(...).pop("params")

On a side note: We are thinking about patterns that avoid having to split the two in the first place. The main reason why you need this is because the optimizer holds the params and the user code holds the state.

[GJBoth]

That's exactly what I do now. I just figured something like

state, params = model.init()

would be much nicer... But yes, not having to split them in the first place would even be better :)