Initialise models from explicit parameters

[nic-barbara]

Add initialisation from explicit model

[nic-barbara]

Working on this at the moment. The current initialisation involves solving an SDP and is not jittable. This means the model.apply() method cannot be jitted, since the setup() method is also called within apply(). An easy work-around for now is to just do this:

explicit = model.params_to_explicit(params)
xnext, y = model.explicit_call(x, u, explicit)

However, this should really be something that we deal with! It might be a fundamental limitation of the way JAX models are constructed, and could be solved by transitioning to NNX as outlined in #23. Unclear, investigate.

[nic-barbara]

Resolved the above by separating out the non-jittable code from the explicit init as a "pre-init" stage. Users now create a REN, then run model.explicit_pre_init() before running model.init() or model.apply(), which fixes everything for now. It's a little ugly, but it'll do just fine for the moment.

[nic-barbara]

Re-opening this issue for future development. The explicit initialisation code has been moved to the dev/explicit-init branch. We need a more streamlined approach for this, but at the moment there's not really much need for it.

[nic-barbara]

Another element to add to this: for contracting RENs, we allow the user to initialise the REN as a linear system. We should also add the option to initialise as almost a linear system of the user's choice. Ian and I discussed this a couple of years ago, suggested sampling from the Wishart distribution. Here's the Julia code from RobustNeuralNetworks.jl that does this:

https://github.com/acfr/RobustNeuralNetworks.jl/blob/main/src/ParameterTypes/contracting_ren.jl#L122