-
Notifications
You must be signed in to change notification settings - Fork 236
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support multi sample guides in Trace_ELBO #1666
Conversation
@@ -116,15 +124,30 @@ class Trace_ELBO(ELBO): | |||
:param vectorize_particles: Whether to use `jax.vmap` to compute ELBOs over the | |||
num_particles-many particles in parallel. If False use `jax.lax.map`. | |||
Defaults to True. | |||
:param multi_sample_guide: Whether to make an assumption that the guide proposes |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why not infer this from isinstance(guide, ...)
? do we expect guides with this attribute apart from a small number of AutoGuide
s we construct?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please feel free to remove this when such auto guides are available.
model_log_density, model_trace = log_density( | ||
replay_model, args, kwargs, params | ||
) | ||
check_model_guide_match(model_trace, guide_trace) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
so we lose these two checks if multi_sample_guide==True
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
we can check the later one. The first one is tricky.
This PR enables running SVI with guides which propose multiple samples.