A PRNGKey can now be passed to the initialisation function for reproducible parameter initialisation when parameters are stochastic e.g., RFFs. Further, the return argument of initialise and fit is a dataclass than bundles up the constituent quantities.