Support vectorized likelihood evaluation

[caoxiaoyue]

I am currently working on a project that aims to enhance the speed of strong lens modeling using JAX. In my particular case, the computation of the likelihood is a relatively time-consuming process, and running the likelihood evaluation on a GPU typically yields faster results compared to using a CPU. By further implementing a vectorized approach for the likelihood evaluation, as described in the UltraNest 3.6.2 documentation's features section, the full potential of the GPU can be harnessed, resulting in a potential speedup of up to 100 times compared to the CPU. I am curious to know if Jaxns can support a similar vectorized likelihood evaluation, similar to the functionality provided by the UltraNest sampler.

[Joshuaalbert]

This should be possible to do. Can you tell me more? What sorts of operations does the likelihood use? Are there while_loops's or scan's? What is the slowest part of the likelihood evaluation, e.g. matrix inverse etc?

And about the model. How many parameters is the model? Is it a complex posterior with sharp edges, etc?

[caoxiaoyue]

Currently, for likelihood computation, I have only implemented a very simple parametric model. The mathematical operations mainly involve internal JAX.numpy functions such as multiplication and arctanh. There are no while_loop or scan operations in my code. The most time-consuming part of calculating the likelihood is the image convolution. Typically, I need to convolve a 200x200 image array with a 21x21 array representing the point spread function. The dimension of our problem is about ~20, by re-parametrizing the model, I believe the likelihood surface does not have sharp edges. I suppose this triangle_plot.pdf can justify this point.

In the past, we used the dynesty software with MCMC step sampling to propose new live points for parameter inference. Recently, we also experimented with the nautilus software (Overview — Nautilus documentation (nautilus-sampler.readthedocs.io)), which employs a neural network to propose new live points. Both dynesty and nautilus samplers have proven to be highly efficient for our lens modeling problem. Additionally, the nautilus sampler supports the vectorized likelihood function that I previously described. However, due to my strong desire to enhance the code's speed, I find the Jax-accelerated Jaxns very appealing. Anyway, if Jaxns can support the vectorized likelihood functionality, I am keen to have a try. Thanks very much for the reply :)