Use module_method for RNNCellBase.initialize_carry about flax HOT 3 CLOSED

I think this is a good idea. I tend to provide my own initialize_carry function, this could be one of

• Constant
• Learnable
• Another model which attempts to predict x0 from previous inputs and outputs

from flax.

@bastings has done some experiments on this issue and they had some insights that I think would be useful to share here. I've assigned this issue to them.

from flax.

Closing this issue as the old nn library is deprecated and we hopefully have a better pattern here in the newer "linen" api:

Just define a non-static initialize_carry e.g. def initialize_carry(self, rng, batch_dims, size, init_fn=zeros): in linen and it will have access to the Module constructor args on self.

Please feel free to re-open (or better make a new issue) pertaining to any perceived deficiency in linen's RNN utils.

from flax.