Make FullyBayesian ABC more flexible again

[CompRhys]

Motivation

In the old ABC we could pass an already initialized pyro_model which allowed for fairly complicated pyro_models to be used in the available framework. The refactor prevents us passing arbitary kwargs to the pyro_model and so is a step back in terms of hackability of the code.

Passing a dict of kwargs seems like the right way to keep things flexible if no longer able to just pass the model. This PR is just a draft as I am not sure yet what else I need to touch.

Have you read the [Contributing Guidelines on pull requests]

Yes

Test Plan

Ensure tests pass, add new tests if needed.

Related PRs

None

[Balandat]

cc @sdaulton

[sdaulton]

This makes sense to me. An alternative to passing pyro_model_kwargs would be to specify the specific args that you'd like to pass, to allow for stricter input checking. What kwargs were you interesting in passing?

[CompRhys]

I am passing in a small neural network to have a fully bayesian deep kernel model and then I hijack the training function to train it. Hence why having the greater flexibility was good

[CompRhys]

@sdaulton what would you suggest are the next steps here? are you happy with the catch-all kwargs and I should work on fixing the tests?

[TobyBoyne]

Just wanted to add that this would be very useful for me! I am intending to use a tree-based GP kernel which we sample with MCMC, which I will convert to a Pyro model, then use as in a FullyBayesian BoTorch model.

[sdaulton]

Sorry for the delay here. I lost track of this issue. A couple thoughts:

1. If we just pass pyro_model_kwargs,

• Then we have little type-checking/validation of the arguments, which is something we aim to avoid in the codebase.
• Then we are still restricted to the pyro_model_class of the specific subclass of AbstractFullyBayesianSingleTaskGP. This wouldn't be able to support something like a tree-based kernel without implementing a new subclass of AbstractFullyBayesianSingleTaskGP (and a new subclass of PyroModel).

2. If we allow passing an instantiated pyro_model, then we'd want to add more subclasses of a base PyroModel in order to be able to better specify the types of PyroModels that are supported by each subclass of AbstractFullyBayesianSingleTaskGP. E.g. a SaasFullyBayesianSingleTaskGP requires an outputscale while a FullyBayesianSingleTaskGP does not allow one.

If one is implementing a custom PyroModel, is there any reason for them not to also implement a custom subclass of AbstractFullyBayesianSingleTaskGP? Curious what the best user experience would look like for using FullyBayesian models