-
Notifications
You must be signed in to change notification settings - Fork 102
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
Enable shared mcmc parameters with tempered smc #694
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -65,16 +65,25 @@ def logprior_fn(x): | |
|
||
hmc_kernel = blackjax.hmc.build_kernel() | ||
hmc_init = blackjax.hmc.init | ||
hmc_parameters = extend_params( | ||
num_particles, | ||
{ | ||
"step_size": 10e-2, | ||
"inverse_mass_matrix": jnp.eye(2), | ||
"num_integration_steps": 50, | ||
}, | ||
hmc_parameters_list = [ | ||
extend_params( | ||
num_particles if extend else 1, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd add a comment highlighting that here you are testing that extending with a copy is the same as having a 1 dimension parameter. Otherwise it may be tricky to get why this is even hapening. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
{ | ||
"step_size": 10e-2, | ||
"inverse_mass_matrix": jnp.eye(2), | ||
"num_integration_steps": 50, | ||
}, | ||
) | ||
for extend in [True, False] | ||
] | ||
hmc_parameters_list.append( | ||
extend_params( | ||
num_particles, {"step_size": 10e-2, "num_integration_steps": 50} | ||
) | ||
| extend_params(num_particles, {"inverse_mass_matrix": jnp.eye(2)}) | ||
) | ||
|
||
for target_ess in [0.5, 0.75]: | ||
for target_ess, hmc_parameters in zip([0.5, 0.5, 0.75], hmc_parameters_list): | ||
tempering = adaptive_tempered_smc( | ||
logprior_fn, | ||
loglikelihood_fn, | ||
|
@@ -115,7 +124,7 @@ def test_fixed_schedule_tempered_smc(self): | |
hmc_init = blackjax.hmc.init | ||
hmc_kernel = blackjax.hmc.build_kernel() | ||
hmc_parameters = extend_params( | ||
100, | ||
1, | ||
{ | ||
"step_size": 10e-2, | ||
"inverse_mass_matrix": jnp.eye(2), | ||
|
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.
calling extend params shouldn't be needed anymore
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.
I'm still using extend_params since it needs to convert to an array and add a leading dimension of length one. I could have added a separate helper but it'd do the same thing as extend_params but just without the jnp.repeat part
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.
I think it is worth to modify
extend_params
into something simple without thejnp.repeat
.Would this work?
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.
That would work, perhaps with potentially a
jnp.asarray()
included. The concern I have with modifyingextend_params
is I'd assume there's still a use case for it with initializing unshared parameters?A current use case is for the remaining tests that are using duplicated parameters for testing the unshared case.
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.
@ciguaran Is there any other use case besides dimension matching initially?
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.
There's none.
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.
I'd assume there's still a use case for it with initializing unshared parameters -> that is a good concern although I think that in most cases we wouldn't initialize by "extending" with copies. We would do it by sampling from some probability distribution. So you can modify it as Junpeng suggested.
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.
ah, that makes sense. thanks