Enabling sampling with intermediates for ExpandedDistribution
#909
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
It involved some elaborate shuffling in sample/batch dimensions that is unneccessary (due to independence in these dimensions) and didn't mitigate the need for a reshape in the end, which achieves the same results all on its own.
Assuming that intermediates are batched similarly to direct samples.
fehiepsi
reviewed
Feb 10, 2021
| shape is correct: This implicitly replaces the interstitial dims | ||
| of size 1 in the original batch_shape of base_dist with those | ||
| in the expanded dims. While it somewhat 'shuffles' over batch | ||
| dimensions, we don't care because they are considered independent.""" |
There was a problem hiding this comment.
Makes sense to me. The change looks great but let me think a bit more to see if I miss something.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Currently,
ExpandedDistributionblocks intermediate values from the base distribution insample_with_intermediates, i.e., intermediate values cannot be retrieved from sampling and traces whenExpandedDistributionand, by extension, theplateprimitive is used.The pull request implements
ExpandedDistribution.sample_with_intermediatesunder the assumption that batch_shape of intermediate values expands in the same way as for the main sampled value.As of yet, there are not additional tests to verify correctness of the intermediates, though.