Find the right boundary between runtimes and programs #13
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.