Skip to content
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.

Sign up for GitHub

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

Tighter conditions for loops in _poisson sampling #583

Merged
merged 2 commits into from
May 5, 2020
Merged

Tighter conditions for loops in _poisson sampling #583

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
74d3e0d
Tighter conditions for loops in _poisson sampling
lumip Apr 30, 2020
96a5bcf
fixing linting issues
lumip May 4, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions numpyro/distributions/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,12 @@ def cond_fn(val):
cond2 = (k < 0) | ((us < 0.013) & (V > us))
cond3 = ((np.log(V) + np.log(invalpha) - np.log(a / (us * us) + b))
<= (-lam + k * loglam - gammaln(k + 1)))
return (~cond1) & (cond2 | (~cond3))

# lax.cond in _poisson_one apparently may still
# execute _poisson_large for small lam:
# additional condition to not iterate if that is the case
cond4 = lam >= 10
return (~cond1) & (cond2 | (~cond3)) & cond4

def body_fn(val):
rng_key, *_ = val
Expand All @@ -221,6 +226,15 @@ def _poisson_small(val):
rng_key, lam = val
enlam = np.exp(-lam)

def cond_fn(val):
cond1 = val[1] > enlam

# lax.cond in _poisson_one apparently may still
# execute _poisson_small for large lam:
# additional condition to not iterate if that is the case
cond2 = lam < 10
return cond1 & cond2

def body_fn(val):
rng_key, prod, k = val
rng_key, key_U = random.split(rng_key)
Expand All @@ -229,7 +243,7 @@ def body_fn(val):
return rng_key, prod, k + 1

init = np.where(lam == 0., 0., -1.)
*_, k = lax.while_loop(lambda val: val[1] > enlam, body_fn, (rng_key, 1., init))
*_, k = lax.while_loop(cond_fn, body_fn, (rng_key, 1., init))
return k


Expand Down