New features for linen api

[machineko]

As I was working with flax for some time on my side project i implemented few thinks that i would like to see maybe upstream or in some sort of addons.

1. Smallest is loss functions (MSE and MAE at the moment) its very basic but still useful to start some sort of flax.loss subpackage so we don't need to write few line of codes in every other repo 🗡️
2. Basic cyclic learning feature and small ported from examples create_learning_rate_scheduler class
3. There is also small method to convert dtypes in frozen params object ( I don't think there is a method for it right now and its useful to convert weights for inference )
4. And last biggest is Transformer (used for FastPitch implementation in flax) based on linen Modules and @struct.dataclass

So I would love to know what from this can be useful in flax main repo?