Can I reuse the sampler to speed up my code?

Hi everyone,

I’m a PyMC newbie and currently using it for performing A/B testing. It’s been great so far, but now I have a usecase in which I have to perform lots of inferences, for example, compute A/B testing results for every day for the last quarter. Due to this, pm.sample() ends up getting called in a for loop which initializes the NUTS sampler in every loop and this seems to take the most amount of time (ranges from ~3 seconds to upto 30 seconds depending on server hardware configuration).

Is there a way to initialize the sampler once and reuse it so that I can save this time? Do I have to “reset” the sampler somehow after performing initialization? I’m not very familiar with the mechanics of how samplers work both at a mathematical and pymc level so I’m unsure how to proceed.

Thanks for the help!

The best may be to use nutpie: GitHub - pymc-devs/nutpie: Python wrapper for nuts-rs

You can compile the model once (the expensive part). And if you define your observations with MutableData, you can do compiled_model.with_data to update the observations. Sampling after this should be pretty fast, as it won’t recompile the logp/dlogp functions.

3 Likes

Thanks for replying! I’ll look into it and respond back.

Worked out great, got ~3x speed up in my workflow. Thank you!

Hi @raghuveer, @ricardoV94.

I am pretty new to the forum. I saw that you can use nutpie for reuse of pymc model. I am wondering if there is a way that I can reuse the sample_blackjax_nuts model as well.

Thanks a lot in advance!

As with nutpie you probably can but you need to interface with the library directly instead of having PyMC do the interface for you