search

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

Save the kernel parameters after warmup #98

Open rlouf opened 3 years ago

rlouf commented 3 years ago

https://github.com/rlouf/mcx/blob/2a2b94801e68d94d86826863eeee80f0b84c390d/mcx/inference/hmc.py#L198-L205

We currently pass the parameters directly to the runtime; While the values of the parameters are passed in the Trace object, it would be convenient to update the Kernel's parameter values.