Avoid nan acceptance probability in SA
#740
Conversation
|
@martinjankowiak Is it better to clip the nonsense value |
|
Anyway, after thinking more about this, I think it is safer to use |
| @@ -172,7 +172,8 @@ def sample_kernel(sa_state, model_args=(), model_kwargs=None): | |||
| log_weights_ = dist.MultivariateNormal(locs_, scale_tril=scales_).log_prob(zs_) + pes_ | |||
| else: | |||
| log_weights_ = dist.Normal(locs_, scales_).log_prob(zs_).sum(-1) + pes_ | |||
| log_weights_ = jnp.where(jnp.isnan(log_weights_), -jnp.inf, log_weights_) | |||
| # mask invalid values (nan, +inf) by -inf | |||
| log_weights_ = jnp.where(jnp.isfinite(log_weights_), log_weights_, -jnp.inf) | |||
There was a problem hiding this comment.
LGTM, but when would this ever be +inf?
There was a problem hiding this comment.
It comes from a diffusion drift model in our forum (where the model has lax.scan). I guess something like 1/0 is involved or precision loss during the loop... (the model is pretty complicated and involves sequential matrix multiplication....).
Currently, when
log_weights_[-1] = inf,will lead to
nanacceptance probability. This PR clipslog_weightsbyfinfo.maxvalue to prevent that issue.cc @awakhloo