Guidance: logdensity function with extra arguments

[johannahaffner]

Hi!

I'm new to Blackjax, and really appreciate the flexibility it promises.
I'm currently thinking about how to implement my first use case.

Can the logdensity function have a signature other than taking a state and spitting out a value?
I understand that this is not possible during the step, at least not in the current implementation of Update_Fn. In the NUTS sampler, it also does not seem to be possible to pass such arguments when building the kernel.
Is it possible for any algorithm in Blackjax?

I have a generalised-least-squares expression to compute the loglikelihood for the parameters of an ODE model, given the data and a parametric variance model. I already use it as an objective function during parameter estimation, and I now want to sample around the estimate. Ideally, I would like to be able to vmap over the data, which represent individual time-series.

From what I have understood so far, the most straightforward thing for me to do would be writing a callable wrapper class for the loglikelihood expression I have and make the other arguments attributes of the wrapper, such that I get an interface that only takes the parameters.

However, this would mean that I need to build the kernel repeatedly (hundreds of times). Will this have a strong impact on the performance of my program?

[johannahaffner]

I found #416, is the use of inner functions the recommended way to create such a wrapper?

In that example, they are being used to create a log density that only depends on the sampled parameter.