General two-step optimization machinery

[BalzaniEdoardo]

Is your feature request related to a problem? Please describe.
Many problems require a two- (or multi)-step optimization procedure:

• EM algorithms
• Negative Binomial joint fit of glm parameters and scale
• Factorized spatio-temporal filters for glm-like models.

We should think about a general infrastructure that can be helpful for all these implementation.

Describe the solution you'd like
Probably, a chain of callables with states, similar to optax chaining mechanism.

func1(param1, state) -> param1, state
func2(param2, state) -> param2, state

Describe alternatives you've considered
Inner/outer loop structure, but that is limited to 2 stages, while the chaining is more versatile.