def test_transform_samples():
aesara.config.on_opt_error = "raise"
np.random.seed(13244)
obs = np.random.normal(10, 2, size=100)
obs_at = aesara.shared(obs, borrow=True, name="obs")
with pm.Model() as model:
a = pm.Uniform("a", -20, 20)
sigma = pm.HalfNormal("sigma")
b = pm.Normal("b", a, sigma=sigma, observed=obs_at)
trace = sample_numpyro_nuts(chains=1, random_seed=1322, keep_untransformed=True)
log_vals = trace.posterior["sigma_log__"].values
trans_vals = trace.posterior["sigma"].values
assert np.allclose(np.exp(log_vals), trans_vals)
assert 8 < trace.posterior["a"].mean() < 11
assert 1.5 < trace.posterior["sigma"].mean() < 2.5
obs_at.set_value(-obs)
with model:
trace = sample_numpyro_nuts(chains=2, random_seed=1322, keep_untransformed=False)
assert -11 < trace.posterior["a"].mean() < -8
assert 1.5 < trace.posterior["sigma"].mean() < 2.5
def test_seeding(chains, random_seed):
sample_kwargs = dict(
tune=100,
draws=5,
chains=chains,
random_seed=random_seed,
)
with pm.Model(rng_seeder=456) as m:
pm.Normal("x", mu=0, sigma=1)
result1 = sample_numpyro_nuts(**sample_kwargs)
with pm.Model(rng_seeder=456) as m:
pm.Normal("x", mu=0, sigma=1)
result2 = sample_numpyro_nuts(**sample_kwargs)
result3 = sample_numpyro_nuts(**sample_kwargs)
assert np.all(result1.posterior["x"] == result2.posterior["x"])
expected_equal_result3 = random_seed is not None
assert np.all(result2.posterior["x"] ==
result3.posterior["x"]) == expected_equal_result3
if chains > 1:
assert np.all(result1.posterior["x"].sel(
chain=0) != result1.posterior["x"].sel(chain=1))
assert np.all(result2.posterior["x"].sel(
chain=0) != result2.posterior["x"].sel(chain=1))
assert np.all(result3.posterior["x"].sel(
chain=0) != result3.posterior["x"].sel(chain=1))
def test_deterministic_samples():
aesara.config.on_opt_error = "raise"
np.random.seed(13244)
obs = np.random.normal(10, 2, size=100)
obs_at = aesara.shared(obs, borrow=True, name="obs")
with pm.Model() as model:
a = pm.Uniform("a", -20, 20)
b = pm.Deterministic("b", a / 2.0)
c = pm.Normal("c", a, sigma=1.0, observed=obs_at)
trace = sample_numpyro_nuts(chains=2, random_seed=1322, keep_untransformed=True)
assert 8 < trace.posterior["a"].mean() < 11
assert np.allclose(trace.posterior["b"].values, trace.posterior["a"].values / 2)