Unified API

[SamDuffield]

This PR provides a unified API for uqlib methods.

The unified API is optax/BlackJAX style with init, update and build functions:

# define log_posterior, params and dataloader

transform = uqlib_method.build(log_posterior, **config_args)

state = transform.init(params)

for batch in dataloader:
       state = transform.update(state, batch)

examples/lightning_autoencoder.py has a full example of how to use the API with pytorch lightning in a way that you can seamlessly swap between sghmc and vi.diag.

[dmitrisaberi]

i don't get what this little bit is doing down here

[SamDuffield]

I just wanted to test that inplace=True indeed modifies the input state in-place, whilst inplace=False leaves the input state untouched.

Definitely this could be improved, although perhaps in a later PR

[dmitrisaberi]

explain this a bit?

[SamDuffield]

Lol yep definitely

[dmitrisaberi]

LGTM, just left some clarifying questions!