tinygp and emcee

[pilchat]

Hi,

I have a code that samples a likelihood function using emcee in a multithreaded framework. the likelihood is computed by tinygp.

Now I have vectorized part of the likelihod function for running on GPUs and wonder about the smartest way to use tinygp in a vectorized flavor. I've been looking for tutorials (not an expert here, unfortunately) but I could only find one of yours using numpyro. Is there any specific reason not to use emcee? More importantly, is there a way to use tinygp in a vectorized way?

Thanks

[dfm]

NumPyro will typically get better sampling performance since it used gradient-based sampling methods, and takes better advantage of JAX's JIT compilation. That being said, there shouldn't be any real reason why you can't use emcee.

Can you say a little more about what you mean by "vectorized" here?

[pilchat]

My emcee sampler samples a vectorized likelihood function. This likelihood functions thus accepts a 2D array and uses the native vectorization of the numpy (or jax.numpy) library. The only problem is that, as far as I understand, the kernel module in tinygp accepts 1D arrays, so I gotta find a way to handle this to run efficiently on a GPU. This point is actually the bottleneck of the entire code, as the GP computation is the most computationally demanding task and I need to optimize it at best.

I have read something about avoiding for loops and using either numpy.vectorize() or map(), but I'm kinda lost among all the possibilities and short on time for proper testings. This is why I'm asking for advice.

Thanks again