Compiling local sampler is slow

[kazewong]

Compiling the sampler for complicated likelihood function seems pretty slow (Testing on a gravitational wave example now.)

I think this is related to defining the main loop inside the sampler.

Experimenting with abstracting that out for performance

[kazewong]

@dfm mentioned using scan to replace fori_loop should help reducing the compilation

[kazewong]

Did some more testing on MALA.py compilation time. I think the main compilation time overhead comes from Jax compiling derivative of the log likelihood multiple times.

As an example, here is the mala kernel for one proposal.

   ...

    key1, key2 = jax.random.split(rng_key)
    proposal = position + dt * d_logpdf(position)
    proposal +=  jnp.sqrt(2 * dt) * jax.random.normal(key1, shape=position.shape)
    ratio = logpdf(proposal) - logpdf(position)
    ratio -= ((position - proposal - dt * d_logpdf(proposal)) ** 2 / (4 * dt)).sum()
    ratio += ((proposal - position - dt * d_logpdf(position)) ** 2 / (4 * dt)).sum()
    proposal_log_prob = logpdf(proposal)

    log_uniform = jnp.log(jax.random.uniform(key2))
    do_accept = log_uniform < ratio

    position = jnp.where(do_accept, proposal, position)
    log_prob = jnp.where(do_accept, proposal_log_prob, log_prob)
    return position, log_prob, do_accept

I tried this example with a gravitational wave likelihood, where one compilation of d_logpdf takes around 70s.
If I have jit the entire kernel, then the compilation time is around 300 seconds.
This seems to indicate Jax does not use the cached version of d_logpdf when jitting the kernel. Even assuming 'proposal' and 'position' somehow trigger recompilation, the compilation time should still be smaller.

Another possibility is Jax unfold the entire mala_kernel computation graph without considering logpdf and d_logpdf are used multiple times hence share the same graph.