Add optimizers to unified API

[SamDuffield]

This PR allows the user to swap out the UQ transforms with optimisation if they like.

Nothing sophisticated just gives a bit more flexibility to the API.

Also open to discussion to whether this is within scope or not.

[dmitrisaberi]

i see how this makes everything more functional, and makes the use of optimizers consistent with the unified api -- and it's very nice in that way -- but i don't get what added flexibility / functionality this actually buys you. could you give me a bit more intuition for this? no problems with code at all though, all LGTM

[SamDuffield]

i see how this makes everything more functional, and makes the use of optimizers consistent with the unified api -- and it's very nice in that way -- but i don't get what added flexibility / functionality this actually buys you. could you give me a bit more intuition for this? no problems with code at all though, all LGTM

I'm not expecting these to replace torch.optim optimisation in general it just adds support in case the user wants to write a script that can run any uqlib method and simple MAP optimisation is likely a good baseline.

Essentially just adds support to swap easily from a UQ method to pure optimisation.

E.g. in the lightning_autoencoder.py example we could have

method, config_args = uqlib.vi.diag, {"optimizer": torchopt.adam(lr=1e-3)}

or

method, config_args = uqlib.optim, {"optimizer": torch.optim.Adam, "lr": 1e-3}

And all of the code below this the same (i.e. method agnostic).

[johnathanchiu]

lgtm