Numpyro: Avoid recompilation in MCMC.run with new data

[neerajprad]

This can be important for saving on compilation cost.

There seems to be some bug in in HMC/NUTS kernels or the MCMC class due to which this doesn't work currently. See. @null-a's snippet for details.

[fehiepsi]

This issue seems to come from how JAX jit works

from jax import jit

class F:
    a = 1

    def sample(self):
        return self.a

f = F()
jit(f.sample)()

f.a = 2
jit(f.sample)()

, which return 1 instead of 2.

[neerajprad]

I thought so too since the compilation overhead on the second call is small, whereas the code as it currently stands should have actually recompiled everything for the second call too. I think this may be a bigger change since this will require us to pipe in args/kwargs through sample fn, if we would like the args to be JITed too.

[neerajprad]

I am marking this as enhancement because not only would we like to return the correct result, but also do so by still not having to recompile everything. This may require changes to a few utilities but hopefully not the MCMC class itself.

[fehiepsi]

not having to recompile everything

I think for new data, we have to recompile sample_fn unless the new data has the same shape. If the data has the same shape, users can jit whole mcmc sampling process (e.g. see this test) at the trade-off of not having progress-bar.

[neerajprad]

Keeping the data shape the same is a reasonable requirement to not having to recompile. I would like to keep this issue open, and ideally solve this by having model arguments be JITable. Then, if the shapes are retained across calls, we don't recompile but if they are, it automatically triggers recompilation.