Skip to content

Commit

Permalink
Add a test for sampling from CustomDist in multiple processes
Browse files Browse the repository at this point in the history
  • Loading branch information
EliasRas committed May 28, 2024
1 parent 19be124 commit 02fcc51
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions tests/smc/test_smc.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,21 @@ def test_unobserved_categorical(self):

assert np.all(np.median(trace["mu"], axis=0) == [1, 2])

def test_parallel_custom(self):
def _logp(value, mu):
return -((value - mu) ** 2)

def _random(mu, rng=None, size=None):
return rng.normal(loc=mu, scale=1, size=size)

def _dist(mu, size=None):
return pm.Normal.dist(mu, 1, size=size)

with pm.Model():
mu = pm.CustomDist("mu", 0, logp=_logp, dist=_dist)
pm.CustomDist("y", mu, logp=_logp, class_name="", random=_random, observed=[1, 2])
pm.sample_smc(draws=6, cores=2)

def test_marginal_likelihood(self):
"""
Verifies that the log marginal likelihood function
Expand Down

0 comments on commit 02fcc51

Please sign in to comment.