Some init or step function orders either fn(state, model) or fn(model, state). I suggest changing everything to fn(state, model). This also involves modifying tests and examples.