Add batch_ndims argument in Predictive and log_likelihood #660
Conversation
| model, data, _ = beta_bernoulli() | ||
| samples = Predictive(model, return_sites=["beta"], num_samples=100)(random.PRNGKey(1)) | ||
| loglik = log_likelihood(model, samples, data) | ||
| samples = Predictive(model, return_sites=["beta"], num_samples=200)(random.PRNGKey(1)) |
There was a problem hiding this comment.
What will happen if we provide batch_ndims=2 here?
There was a problem hiding this comment.
I believe the output will have shape (1, 200, ...) but let me test it.
There was a problem hiding this comment.
That's true, I just add a test for that in test_prior_predictive.
numpyro/distributions/transforms.py
Outdated
| warnings.warn(DeprecationWarning, | ||
| "MultivariateAffineTransform is renamed to LowerCholeskyAffine.") | ||
| warnings.warn("MultivariateAffineTransform is renamed to LowerCholeskyAffine.", | ||
| DeprecationWarning) |
There was a problem hiding this comment.
Sorry, I might have missed this earlier.DeprecationgWarning will be swallowed by application code. We should use FutureWarning instead. See https://docs.python.org/3/library/warnings.html#warning-categories.
There was a problem hiding this comment.
Thanks! This explains why I didn't see the warnings when I used predictive.get_samples(...). :D
| @@ -32,8 +32,8 @@ def init_to_median(site=None, num_samples=15): | |||
|
|
|||
|
|
|||
| def init_to_prior(site=None): | |||
| warnings.warn(DeprecationWarning, | |||
| "`init_to_prior` strategy is renamed to `init_to_sample`.") | |||
| warnings.warn("`init_to_prior` strategy is renamed to `init_to_sample`.", | |||
There was a problem hiding this comment.
same here about DeprecationWarning.
Addresses #645 and #639.
This PR also adds a utility
soft_vmapto avoid duplicated code in Predictive and log_likelihood. It might also be useful for other purposes. @martinjankowiak you might findsoft_vmap, which is adapted from your code, useful for your experiment. :D