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.

2 Likes

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

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