Losses currently cannot use random state, while its not that common its sometime needed and it should be supported. There are two ways of going about this:
Let Losses optionally accept a next_key: KeySeq argument, no modifications to Loss needed but now Model need to have a next_key: KeySeq field (not that bad).
Let Loss inherit from treex.Treex or even just treeo.Tree so the user can just create its own random state. Downside is that losses are potentially stateful and we would need watch out for the implications when used inside a Metric e.g. Losses and friends.
Losses currently cannot use random state, while its not that common its sometime needed and it should be supported. There are two ways of going about this:
next_key: KeySeqargument, no modifications toLossneeded but nowModelneed to have anext_key: KeySeqfield (not that bad).Lossinherit fromtreex.Treexor even justtreeo.Treeso the user can just create its own random state. Downside is that losses are potentially stateful and we would need watch out for the implications when used inside aMetrice.g.Lossesand friends.