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.
rlouf
commented
4 years ago
I am not sure what runtimes and programs' respective responsibilities are yet. This PR is an experiment to find the right API.
Requirements
inferencemodule;Slow initialisation
The initialisation is unacceptably slow. It is caused by repeatedly calling
ravel_pytreefor each chain. Luckily for us,ravel_pytreearranges 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_forwardto obtained flattened position, and applyravel_pytreeto one dictionary to get the unraveling function.This PR follows issue #12.