implements an option to provide kernels with a logprob based on subsets - #296

merged main into this PR.

@jobrachem, could we discuss the design soonish? I'd like to adjust the other kernels and then merge this.

Hey @wiep, the implementation looks solid and straight-forward to use. I have some general thoughts:

• Currently, the implementation requires changes to the interface of all kernels, adding an init argument that is not (and probably never will be) compatible with other interfaces than the LieselInterface.
• I am wondering whether the problem could maybe also be solved on the level of the model interface instead. For the kernels, it would then look like any other interface, no change required. To allow different kernels to use different interfaces would require some changes to the builder - but maybe that would work out. We could, for example, let the Interface.log_prob method accept **kwargs, and then allow users to pass such kwargs in EngineBuilder.add_kernel. Then, if the LieselInterface.log_prob receives a log_prob_node_name kwarg, it returns model_state[log_prob_node_name].value. Such a design would not require bespoke changes to any kernels (except for passing on the kwargs) and open up similar (and more) modifications also for other interfaces. Example:

eb.add_kernel(kernel, log_prob_kwargs={"log_prob_node_name": "my_subset_logp"})

What do you think?

to clarify, you would recommend leaving the kernels unchanged but basically giving them a different model interface. this model interface will be created by the engine builder and for the kernel it is transparent and it would not know that it operates on a different model than other kernels?

in general that sounds like a solid idea. However, we need to consider:

• does this work with the inference_spec of nodes?
• does it separate kernel creation and the model adjustment too much, i.e. are they at different parts in the code?

thinking about the MCMCSpec it would basically require us to create a new dictionary of kwargs that are passed to the eb.add_kernel method.

I like that we would not need to touch the kernels at all and just give them a model interface with a different log_prob method. But is MCMCSpec then too complicated?

to clarify, you would recommend leaving the kernels unchanged but basically giving them a different model interface. this model interface will be created by the engine builder and for the kernel it is transparent and it would not know that it operates on a different model than other kernels?

Yes, that is one of two ways that I see for going through the interfaces.

The other one would be to provide the possibility of using the same interface in all kernels, but allow different kernels to pass different keyword arguments to the log_prob method of the interface. This would have the downside that we would indeed need changes to the kernels (for letting them pass on the kwargs).

in general that sounds like a solid idea. However, we need to consider:

• does this work with the inference_spec of nodes?
• does it separate kernel creation and the model adjustment too much, i.e. are they at different parts in the code?

I suppose that would indeed need some thought.

thinking about the MCMCSpec it would basically require us to create a new dictionary of kwargs that are passed to the eb.add_kernel method.

That could be it.

I like that we would not need to touch the kernels at all and just give them a model interface with a different log_prob method.

But is MCMCSpec then too complicated?

I guess it is a question of tradeoffs. The code itself wouldn't be too bad I think. What's your intuition?

i think it would be great to keep the kernels simple and not add additional complexity that needs to be repeated for every kernel implementation.

I agree