rlouf / mcx

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.
https://rlouf.github.io/mcx
Apache License 2.0
325 stars 17 forks source link

Merge batch and iterative samplers #38

Closed rlouf closed 3 years ago

rlouf commented 3 years ago

Since batch and iterative sampling share the same API and philosophy it would be preferable to combine both execution models in the same object. We want to be able to do things like:

sampler = mcx.sampler(...)
for sample in sampler:
    pass

or

sampler = mcx.sampler(...)
iter(sampler)
sample = next(sampler)

The first behavior is already implemented, the sampler is an iterable.