I am not sure what runtimes and programs' respective responsibilities are yet. This PR is an experiment to find the right API.
Requirements
The programs need not be aware of the models so they can be used with custom loglikelihoods. As such they can be moved to the inference module;
It should be simple to use meta-algorithms like HMC Swindles and Replica exchange with an existing program;
Runtimes are program-agnostic; they communicate via a unified API;
One should be able to choose the integrator used in HMC;
Warmup adapts to the parameters passed to the program;
Runtimes come with reasonable defaults (dynamicHMC or empiricalHMC), but can easily be used with any algorithm.
Slow initialisation
The initialisation is unacceptably slow. It is caused by repeatedly calling ravel_pytree for each chain. Luckily for us, ravel_pytree arranges the values of the dictionary in a deterministic way, sorting them by key. To fix the performance, we stack the arrays in the dictionary returned by sample_forward to obtained flattened position, and apply ravel_pytreeto one dictionary to get the unraveling function.
I am not sure what runtimes and programs' respective responsibilities are yet. This PR is an experiment to find the right API.
Requirements
inference
module;Slow initialisation
The initialisation is unacceptably slow. It is caused by repeatedly calling
ravel_pytree
for each chain. Luckily for us,ravel_pytree
arranges the values of the dictionary in a deterministic way, sorting them by key. To fix the performance, we stack the arrays in the dictionary returned bysample_forward
to obtained flattened position, and applyravel_pytree
to one dictionary to get the unraveling function.This PR follows issue #12.