def test_logistic_regression():
from tensorflow_probability.substrates.jax import distributions as tfd
N, dim = 3000, 3
num_warmup, num_samples = (1000, 1000)
data = random.normal(random.PRNGKey(0), (N, dim))
true_coefs = jnp.arange(1.0, dim + 1.0)
logits = jnp.sum(true_coefs * data, axis=-1)
labels = tfd.Bernoulli(logits=logits).sample(seed=random.PRNGKey(1))
def model(labels):
coefs = numpyro.sample("coefs",
tfd.Normal(jnp.zeros(dim), jnp.ones(dim)))
logits = numpyro.deterministic("logits", jnp.sum(coefs * data,
axis=-1))
return numpyro.sample("obs", tfd.Bernoulli(logits=logits), obs=labels)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
mcmc.run(random.PRNGKey(2), labels)
mcmc.print_summary()
samples = mcmc.get_samples()
assert samples["logits"].shape == (num_samples, N)
expected_coefs = jnp.array([0.97, 2.05, 3.18])
assert_allclose(jnp.mean(samples["coefs"], 0), expected_coefs, atol=0.22)
def test_beta_bernoulli():
from tensorflow_probability.substrates.jax import distributions as tfd
num_warmup, num_samples = (500, 2000)
def model(data):
alpha = jnp.array([1.1, 1.1])
beta = jnp.array([1.1, 1.1])
p_latent = numpyro.sample("p_latent", tfd.Beta(alpha, beta))
numpyro.sample("obs", tfd.Bernoulli(p_latent), obs=data)
return p_latent
true_probs = jnp.array([0.9, 0.1])
data = tfd.Bernoulli(true_probs).sample(seed=random.PRNGKey(1),
sample_shape=(1000, 2))
kernel = NUTS(model=model, trajectory_length=0.1)
mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
mcmc.run(random.PRNGKey(2), data)
mcmc.print_summary()
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples["p_latent"], 0), true_probs, atol=0.05)
def model(data):
alpha = jnp.array([1.1, 1.1])
beta = jnp.array([1.1, 1.1])
p_latent = numpyro.sample("p_latent", tfd.Beta(alpha, beta))
numpyro.sample("obs", tfd.Bernoulli(p_latent), obs=data)
return p_latent
def model(labels):
coefs = numpyro.sample("coefs",
tfd.Normal(jnp.zeros(dim), jnp.ones(dim)))
logits = numpyro.deterministic("logits", jnp.sum(coefs * data,
axis=-1))
return numpyro.sample("obs", tfd.Bernoulli(logits=logits), obs=labels)
def model(y):
theta = numpyro.sample("p", tfd.Beta(1, 1))
with numpyro.plate("plate", y.size):
numpyro.sample("y", tfd.Bernoulli(probs=theta), obs=y)