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

Improve separation of concerns between executors and algorithms #12

Closed rlouf closed 4 years ago

rlouf commented 4 years ago

Algorithms shouldn’t have to be aware of the existence of the model; everything model-related should happen in the executors, and can be shared between them. Algorithm definitions can be moved back to inference.

This way, algorithms could also be used with custom logpdfs.

hmc = HMC(
    step_size=0.01,
    num_integration_steps=10,
    integrator=mclachlan_integrator  # it should be easy to switch 
)

# `sample`alone should be good default. `using` is for people who want to use alternative algorithms  
sampler = mcx.sample(model, num_chains=4).using(hmc)
rlouf commented 4 years ago

PR #13 addresses this issue, closing.