Python jax.numpy.mean() Examples
The following are 30
code examples of jax.numpy.mean().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
jax.numpy
, or try the search function
.
![](https://www.programcreek.com/common/static/images/search.png)
Example #1
Source File: test_examples.py From jaxnet with Apache License 2.0 | 6 votes |
def test_readme(): net = Sequential(Dense(1024), relu, Dense(1024), relu, Dense(4), log_softmax) @parametrized def loss(inputs, targets): return -jnp.mean(net(inputs) * targets) def next_batch(): return jnp.zeros((3, 784)), jnp.zeros((3, 4)) params = loss.init_parameters(*next_batch(), key=PRNGKey(0)) print(params.sequential.dense2.bias) # [-0.01101029, -0.00749435, -0.00952365, 0.00493979] assert jnp.allclose([-0.01101029, -0.00749435, -0.00952365, 0.00493979], params.sequential.dense2.bias) out = loss.apply(params, *next_batch()) assert () == out.shape out_ = loss.apply(params, *next_batch(), jit=True) assert out.shape == out_.shape
Example #2
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_improper_normal(): true_coef = 0.9 def model(data): alpha = numpyro.sample('alpha', dist.Uniform(0, 1)) with numpyro.handlers.reparam(config={'loc': TransformReparam()}): loc = numpyro.sample('loc', dist.TransformedDistribution( dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha))) numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data) data = true_coef + random.normal(random.PRNGKey(0), (1000,)) kernel = NUTS(model=model) mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000) mcmc.run(random.PRNGKey(0), data) samples = mcmc.get_samples() assert_allclose(jnp.mean(samples['loc'], 0), true_coef, atol=0.05)
Example #3
Source File: bnn.py From numpyro with Apache License 2.0 | 6 votes |
def get_data(N=50, D_X=3, sigma_obs=0.05, N_test=500): D_Y = 1 # create 1d outputs np.random.seed(0) X = jnp.linspace(-1, 1, N) X = jnp.power(X[:, np.newaxis], jnp.arange(D_X)) W = 0.5 * np.random.randn(D_X) Y = jnp.dot(X, W) + 0.5 * jnp.power(0.5 + X[:, 1], 2.0) * jnp.sin(4.0 * X[:, 1]) Y += sigma_obs * np.random.randn(N) Y = Y[:, np.newaxis] Y -= jnp.mean(Y) Y /= jnp.std(Y) assert X.shape == (N, D_X) assert Y.shape == (N, D_Y) X_test = jnp.linspace(-1.3, 1.3, N_test) X_test = jnp.power(X_test[:, np.newaxis], jnp.arange(D_X)) return X, Y, X_test
Example #4
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_beta_bernoulli_x64(kernel_cls): warmup_steps, num_samples = (100000, 100000) if kernel_cls is SA else (500, 20000) def model(data): alpha = jnp.array([1.1, 1.1]) beta = jnp.array([1.1, 1.1]) p_latent = numpyro.sample('p_latent', dist.Beta(alpha, beta)) numpyro.sample('obs', dist.Bernoulli(p_latent), obs=data) return p_latent true_probs = jnp.array([0.9, 0.1]) data = dist.Bernoulli(true_probs).sample(random.PRNGKey(1), (1000, 2)) if kernel_cls is SA: kernel = SA(model=model) else: kernel = kernel_cls(model=model, trajectory_length=0.1) mcmc = MCMC(kernel, num_warmup=warmup_steps, num_samples=num_samples, progress_bar=False) 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) if 'JAX_ENABLE_X64' in os.environ: assert samples['p_latent'].dtype == jnp.float64
Example #5
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_dirichlet_categorical_x64(kernel_cls, dense_mass): warmup_steps, num_samples = 100, 20000 def model(data): concentration = jnp.array([1.0, 1.0, 1.0]) p_latent = numpyro.sample('p_latent', dist.Dirichlet(concentration)) numpyro.sample('obs', dist.Categorical(p_latent), obs=data) return p_latent true_probs = jnp.array([0.1, 0.6, 0.3]) data = dist.Categorical(true_probs).sample(random.PRNGKey(1), (2000,)) kernel = kernel_cls(model, trajectory_length=1., dense_mass=dense_mass) mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False) mcmc.run(random.PRNGKey(2), data) samples = mcmc.get_samples() assert_allclose(jnp.mean(samples['p_latent'], 0), true_probs, atol=0.02) if 'JAX_ENABLE_X64' in os.environ: assert samples['p_latent'].dtype == jnp.float64
Example #6
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_correlated_mvn(): # This requires dense mass matrix estimation. D = 5 warmup_steps, num_samples = 5000, 8000 true_mean = 0. a = jnp.tril(0.5 * jnp.fliplr(jnp.eye(D)) + 0.1 * jnp.exp(random.normal(random.PRNGKey(0), shape=(D, D)))) true_cov = jnp.dot(a, a.T) true_prec = jnp.linalg.inv(true_cov) def potential_fn(z): return 0.5 * jnp.dot(z.T, jnp.dot(true_prec, z)) init_params = jnp.zeros(D) kernel = NUTS(potential_fn=potential_fn, dense_mass=True) mcmc = MCMC(kernel, warmup_steps, num_samples) mcmc.run(random.PRNGKey(0), init_params=init_params) samples = mcmc.get_samples() assert_allclose(jnp.mean(samples), true_mean, atol=0.02) assert np.sum(np.abs(np.cov(samples.T) - true_cov)) / D**2 < 0.02
Example #7
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_binomial_stable_x64(with_logits): # Ref: https://github.com/pyro-ppl/pyro/issues/1706 warmup_steps, num_samples = 200, 200 def model(data): p = numpyro.sample('p', dist.Beta(1., 1.)) if with_logits: logits = logit(p) numpyro.sample('obs', dist.Binomial(data['n'], logits=logits), obs=data['x']) else: numpyro.sample('obs', dist.Binomial(data['n'], probs=p), obs=data['x']) data = {'n': 5000000, 'x': 3849} kernel = NUTS(model=model) mcmc = MCMC(kernel, warmup_steps, num_samples) mcmc.run(random.PRNGKey(2), data) samples = mcmc.get_samples() assert_allclose(jnp.mean(samples['p'], 0), data['x'] / data['n'], rtol=0.05) if 'JAX_ENABLE_X64' in os.environ: assert samples['p'].dtype == jnp.float64
Example #8
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_mcmc_progbar(): true_mean, true_std = 1., 2. num_warmup, num_samples = 10, 10 def model(data): mean = numpyro.sample('mean', dist.Normal(0, 1).mask(False)) std = numpyro.sample('std', dist.LogNormal(0, 1).mask(False)) return numpyro.sample('obs', dist.Normal(mean, std), obs=data) data = dist.Normal(true_mean, true_std).sample(random.PRNGKey(1), (2000,)) kernel = NUTS(model=model) mcmc = MCMC(kernel, num_warmup, num_samples) mcmc.warmup(random.PRNGKey(2), data) mcmc.run(random.PRNGKey(3), data) mcmc1 = MCMC(kernel, num_warmup, num_samples, progress_bar=False) mcmc1.run(random.PRNGKey(2), data) with pytest.raises(AssertionError): check_close(mcmc1.get_samples(), mcmc.get_samples(), atol=1e-4, rtol=1e-4) mcmc1.warmup(random.PRNGKey(2), data) mcmc1.run(random.PRNGKey(3), data) check_close(mcmc1.get_samples(), mcmc.get_samples(), atol=1e-4, rtol=1e-4) check_close(mcmc1._warmup_state, mcmc._warmup_state, atol=1e-4, rtol=1e-4)
Example #9
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_chain(use_init_params, chain_method): N, dim = 3000, 3 num_chains = 2 num_warmup, num_samples = 5000, 5000 data = random.normal(random.PRNGKey(0), (N, dim)) true_coefs = jnp.arange(1., dim + 1.) logits = jnp.sum(true_coefs * data, axis=-1) labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1)) def model(labels): coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim))) logits = jnp.sum(coefs * data, axis=-1) return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels) kernel = NUTS(model=model) mcmc = MCMC(kernel, num_warmup, num_samples, num_chains=num_chains) mcmc.chain_method = chain_method init_params = None if not use_init_params else \ {'coefs': jnp.tile(jnp.ones(dim), num_chains).reshape(num_chains, dim)} mcmc.run(random.PRNGKey(2), labels, init_params=init_params) samples_flat = mcmc.get_samples() assert samples_flat['coefs'].shape[0] == num_chains * num_samples samples = mcmc.get_samples(group_by_chain=True) assert samples['coefs'].shape[:2] == (num_chains, num_samples) assert_allclose(jnp.mean(samples_flat['coefs'], 0), true_coefs, atol=0.21)
Example #10
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_unnormalized_normal_x64(kernel_cls, dense_mass): true_mean, true_std = 1., 0.5 warmup_steps, num_samples = (100000, 100000) if kernel_cls is SA else (1000, 8000) def potential_fn(z): return 0.5 * jnp.sum(((z - true_mean) / true_std) ** 2) init_params = jnp.array(0.) if kernel_cls is SA: kernel = SA(potential_fn=potential_fn, dense_mass=dense_mass) else: kernel = kernel_cls(potential_fn=potential_fn, trajectory_length=8, dense_mass=dense_mass) mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False) mcmc.run(random.PRNGKey(0), init_params=init_params) mcmc.print_summary() hmc_states = mcmc.get_samples() assert_allclose(jnp.mean(hmc_states), true_mean, rtol=0.07) assert_allclose(jnp.std(hmc_states), true_std, rtol=0.07) if 'JAX_ENABLE_X64' in os.environ: assert hmc_states.dtype == jnp.float64
Example #11
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_functional_beta_bernoulli_x64(algo): warmup_steps, num_samples = 500, 20000 def model(data): alpha = jnp.array([1.1, 1.1]) beta = jnp.array([1.1, 1.1]) p_latent = numpyro.sample('p_latent', dist.Beta(alpha, beta)) numpyro.sample('obs', dist.Bernoulli(p_latent), obs=data) return p_latent true_probs = jnp.array([0.9, 0.1]) data = dist.Bernoulli(true_probs).sample(random.PRNGKey(1), (1000, 2)) init_params, potential_fn, constrain_fn, _ = initialize_model(random.PRNGKey(2), model, model_args=(data,)) init_kernel, sample_kernel = hmc(potential_fn, algo=algo) hmc_state = init_kernel(init_params, trajectory_length=1., num_warmup=warmup_steps) samples = fori_collect(0, num_samples, sample_kernel, hmc_state, transform=lambda x: constrain_fn(x.z)) assert_allclose(jnp.mean(samples['p_latent'], 0), true_probs, atol=0.05) if 'JAX_ENABLE_X64' in os.environ: assert samples['p_latent'].dtype == jnp.float64
Example #12
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_functional_map(algo, map_fn): if map_fn is pmap and xla_bridge.device_count() == 1: pytest.skip('pmap test requires device_count greater than 1.') true_mean, true_std = 1., 2. warmup_steps, num_samples = 1000, 8000 def potential_fn(z): return 0.5 * jnp.sum(((z - true_mean) / true_std) ** 2) init_kernel, sample_kernel = hmc(potential_fn, algo=algo) init_params = jnp.array([0., -1.]) rng_keys = random.split(random.PRNGKey(0), 2) init_kernel_map = map_fn(lambda init_param, rng_key: init_kernel( init_param, trajectory_length=9, num_warmup=warmup_steps, rng_key=rng_key)) init_states = init_kernel_map(init_params, rng_keys) fori_collect_map = map_fn(lambda hmc_state: fori_collect(0, num_samples, sample_kernel, hmc_state, transform=lambda x: x.z, progbar=False)) chain_samples = fori_collect_map(init_states) assert_allclose(jnp.mean(chain_samples, axis=1), jnp.repeat(true_mean, 2), rtol=0.06) assert_allclose(jnp.std(chain_samples, axis=1), jnp.repeat(true_std, 2), rtol=0.06)
Example #13
Source File: test_examples.py From jaxnet with Apache License 2.0 | 6 votes |
def test_mnist_vae(): @parametrized def encode(input): input = Sequential(Dense(5), relu, Dense(5), relu)(input) mean = Dense(10)(input) variance = Sequential(Dense(10), softplus)(input) return mean, variance decode = Sequential(Dense(5), relu, Dense(5), relu, Dense(5 * 5)) @parametrized def elbo(key, images): mu_z, sigmasq_z = encode(images) logits_x = decode(gaussian_sample(key, mu_z, sigmasq_z)) return bernoulli_logpdf(logits_x, images) - gaussian_kl(mu_z, sigmasq_z) params = elbo.init_parameters(PRNGKey(0), jnp.zeros((32, 5 * 5)), key=PRNGKey(0)) assert (5, 10) == params.encode.sequential1.dense.kernel.shape
Example #14
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_get_proposal_loc_and_scale(dense_mass): N = 10 dim = 3 samples = random.normal(random.PRNGKey(0), (N, dim)) loc = jnp.mean(samples[:-1], 0) if dense_mass: scale = jnp.linalg.cholesky(jnp.cov(samples[:-1], rowvar=False, bias=True)) else: scale = jnp.std(samples[:-1], 0) actual_loc, actual_scale = _get_proposal_loc_and_scale(samples[:-1], loc, scale, samples[-1]) expected_loc, expected_scale = [], [] for i in range(N - 1): samples_i = np.delete(samples, i, axis=0) expected_loc.append(jnp.mean(samples_i, 0)) if dense_mass: expected_scale.append(jnp.linalg.cholesky(jnp.cov(samples_i, rowvar=False, bias=True))) else: expected_scale.append(jnp.std(samples_i, 0)) expected_loc = jnp.stack(expected_loc) expected_scale = jnp.stack(expected_scale) assert_allclose(actual_loc, expected_loc, rtol=1e-4) assert_allclose(actual_scale, expected_scale, atol=1e-6, rtol=0.05)
Example #15
Source File: test_infer_util.py From numpyro with Apache License 2.0 | 6 votes |
def test_predictive(parallel): model, data, true_probs = beta_bernoulli() mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100) mcmc.run(random.PRNGKey(0), data) samples = mcmc.get_samples() predictive = Predictive(model, samples, parallel=parallel) predictive_samples = predictive(random.PRNGKey(1)) assert predictive_samples.keys() == {"beta_sq", "obs"} predictive.return_sites = ["beta", "beta_sq", "obs"] predictive_samples = predictive(random.PRNGKey(1)) # check shapes assert predictive_samples["beta"].shape == (100,) + true_probs.shape assert predictive_samples["beta_sq"].shape == (100,) + true_probs.shape assert predictive_samples["obs"].shape == (100,) + data.shape # check sample mean assert_allclose(predictive_samples["obs"].reshape((-1,) + true_probs.shape).mean(0), true_probs, rtol=0.1)
Example #16
Source File: pixelcnn.py From jaxnet with Apache License 2.0 | 6 votes |
def ConvOrConvTranspose(out_chan, filter_shape=(3, 3), strides=None, padding='SAME', init_scale=1., transpose=False): strides = strides or (1,) * len(filter_shape) def apply(inputs, V, g, b): V = g * _l2_normalize(V, (0, 1, 2)) return (lax.conv_transpose if transpose else _conv)(inputs, V, strides, padding) - b @parametrized def conv_or_conv_transpose(inputs): V = parameter(filter_shape + (inputs.shape[-1], out_chan), normal(.05), 'V') example_out = apply(inputs, V=V, g=jnp.ones(out_chan), b=jnp.zeros(out_chan)) # TODO remove need for `.aval.val` when capturing variables in initializer function: g = Parameter(lambda key: init_scale / jnp.sqrt(jnp.var(example_out.aval.val, (0, 1, 2)) + 1e-10), 'g')() b = Parameter(lambda key: jnp.mean(example_out.aval.val, (0, 1, 2)) * g.aval.val, 'b')() return apply(inputs, V, b, g) return conv_or_conv_transpose
Example #17
Source File: ucbadmit.py From numpyro with Apache License 2.0 | 6 votes |
def main(args): _, fetch_train = load_dataset(UCBADMIT, split='train', shuffle=False) dept, male, applications, admit = fetch_train() rng_key, rng_key_predict = random.split(random.PRNGKey(1)) zs = run_inference(dept, male, applications, admit, rng_key, args) pred_probs = Predictive(glmm, zs)(rng_key_predict, dept, male, applications)['probs'] header = '=' * 30 + 'glmm - TRAIN' + '=' * 30 print_results(header, pred_probs, dept, male, admit / applications) # make plots fig, ax = plt.subplots(1, 1) ax.plot(range(1, 13), admit / applications, "o", ms=7, label="actual rate") ax.errorbar(range(1, 13), jnp.mean(pred_probs, 0), jnp.std(pred_probs, 0), fmt="o", c="k", mfc="none", ms=7, elinewidth=1, label=r"mean $\pm$ std") ax.plot(range(1, 13), jnp.percentile(pred_probs, 5, 0), "k+") ax.plot(range(1, 13), jnp.percentile(pred_probs, 95, 0), "k+") ax.set(xlabel="cases", ylabel="admit rate", title="Posterior Predictive Check with 90% CI") ax.legend() plt.savefig("ucbadmit_plot.pdf") plt.tight_layout()
Example #18
Source File: test_infer_util.py From numpyro with Apache License 2.0 | 6 votes |
def test_predictive_with_improper(): true_coef = 0.9 def model(data): alpha = numpyro.sample('alpha', dist.Uniform(0, 1)) with handlers.reparam(config={'loc': TransformReparam()}): loc = numpyro.sample('loc', dist.TransformedDistribution( dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha))) numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data) data = true_coef + random.normal(random.PRNGKey(0), (1000,)) kernel = NUTS(model=model) mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000) mcmc.run(random.PRNGKey(0), data) samples = mcmc.get_samples() obs_pred = Predictive(model, samples)(random.PRNGKey(1), data=None)["obs"] assert_allclose(jnp.mean(obs_pred), true_coef, atol=0.05)
Example #19
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 6 votes |
def get_data(N=20, S=2, P=10, sigma_obs=0.05): assert S < P and P > 1 and S > 0 np.random.seed(0) X = np.random.randn(N, P) # generate S coefficients with non-negligible magnitude W = 0.5 + 2.5 * np.random.rand(S) # generate data using the S coefficients and a single pairwise interaction Y = np.sum(X[:, 0:S] * W, axis=-1) + X[:, 0] * X[:, 1] + sigma_obs * np.random.randn(N) Y -= jnp.mean(Y) Y_std = jnp.std(Y) assert X.shape == (N, P) assert Y.shape == (N,) return X, Y / Y_std, W / Y_std, 1.0 / Y_std # Helper function for analyzing the posterior statistics for coefficient theta_i
Example #20
Source File: modules.py From jaxnet with Apache License 2.0 | 6 votes |
def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True, beta_init=zeros, gamma_init=ones): """Layer construction function for a batch normalization layer.""" axis = (axis,) if np.isscalar(axis) else axis @parametrized def batch_norm(x): ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x))) mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True) z = (x - mean) / np.sqrt(var + epsilon) shape = tuple(d for i, d in enumerate(x.shape) if i not in axis) scaled = z * parameter(shape, gamma_init, 'gamma')[ed] if scale else z return scaled + parameter(shape, beta_init, 'beta')[ed] if center else scaled return batch_norm
Example #21
Source File: test_hmc_util.py From numpyro with Apache License 2.0 | 6 votes |
def test_gaussian_subposterior(method, diagonal): D = 10 n_samples = 10000 n_draws = 9000 n_subs = 8 mean = jnp.arange(D) cov = jnp.ones((D, D)) * 0.9 + jnp.identity(D) * 0.1 subcov = n_subs * cov # subposterior's covariance subposteriors = list(dist.MultivariateNormal(mean, subcov).sample( random.PRNGKey(1), (n_subs, n_samples))) draws = method(subposteriors, n_draws, diagonal=diagonal) assert draws.shape == (n_draws, D) assert_allclose(jnp.mean(draws, axis=0), mean, atol=0.03) if diagonal: assert_allclose(jnp.var(draws, axis=0), jnp.diag(cov), atol=0.05) else: assert_allclose(jnp.cov(draws.T), cov, atol=0.05)
Example #22
Source File: test_autoguide.py From numpyro with Apache License 2.0 | 6 votes |
def test_beta_bernoulli(auto_class): data = jnp.array([[1.0] * 8 + [0.0] * 2, [1.0] * 4 + [0.0] * 6]).T def model(data): f = numpyro.sample('beta', dist.Beta(jnp.ones(2), jnp.ones(2))) numpyro.sample('obs', dist.Bernoulli(f), obs=data) adam = optim.Adam(0.01) guide = auto_class(model, init_strategy=init_strategy) svi = SVI(model, guide, adam, ELBO()) svi_state = svi.init(random.PRNGKey(1), data) def body_fn(i, val): svi_state, loss = svi.update(val, data) return svi_state svi_state = fori_loop(0, 3000, body_fn, svi_state) params = svi.get_params(svi_state) true_coefs = (jnp.sum(data, axis=0) + 1) / (data.shape[0] + 2) # test .sample_posterior method posterior_samples = guide.sample_posterior(random.PRNGKey(1), params, sample_shape=(1000,)) assert_allclose(jnp.mean(posterior_samples['beta'], 0), true_coefs, atol=0.05)
Example #23
Source File: test_reparam.py From numpyro with Apache License 2.0 | 5 votes |
def get_moments(x): assert (x > 0).all() x = jnp.log(x) m1 = jnp.mean(x, axis=0) x = x - m1 xx = x * x xxx = x * xx xxxx = xx * xx m2 = jnp.mean(xx, axis=0) m3 = jnp.mean(xxx, axis=0) / m2 ** 1.5 m4 = jnp.mean(xxxx, axis=0) / m2 ** 2 return jnp.stack([m1, m2, m3, m4])
Example #24
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def _mvn_to_scipy(loc, cov, prec, tril): jax_dist = dist.MultivariateNormal(loc, cov, prec, tril) mean = jax_dist.mean cov = jax_dist.covariance_matrix return osp.multivariate_normal(mean=mean, cov=cov)
Example #25
Source File: test_distributions_util.py From numpyro with Apache License 2.0 | 5 votes |
def test_binomial_mean(n, p): samples = binomial(random.PRNGKey(1), p, n, shape=(100, 100)) expected_mean = n * p assert_allclose(jnp.mean(samples), expected_mean, rtol=0.05)
Example #26
Source File: gp.py From numpyro with Apache License 2.0 | 5 votes |
def main(args): X, Y, X_test = get_data(N=args.num_data) # do inference rng_key, rng_key_predict = random.split(random.PRNGKey(0)) samples = run_inference(model, args, rng_key, X, Y) # do prediction vmap_args = (random.split(rng_key_predict, args.num_samples * args.num_chains), samples['kernel_var'], samples['kernel_length'], samples['kernel_noise']) means, predictions = vmap(lambda rng_key, var, length, noise: predict(rng_key, X, Y, X_test, var, length, noise))(*vmap_args) mean_prediction = np.mean(means, axis=0) percentiles = np.percentile(predictions, [5.0, 95.0], axis=0) # make plots fig, ax = plt.subplots(1, 1) # plot training data ax.plot(X, Y, 'kx') # plot 90% confidence level of predictions ax.fill_between(X_test, percentiles[0, :], percentiles[1, :], color='lightblue') # plot mean prediction ax.plot(X_test, mean_prediction, 'blue', ls='solid', lw=2.0) ax.set(xlabel="X", ylabel="Y", title="Mean predictions with 90% CI") plt.savefig("gp_plot.pdf") plt.tight_layout()
Example #27
Source File: test_autoguide.py From numpyro with Apache License 2.0 | 5 votes |
def test_logistic_regression(auto_class): N, dim = 3000, 3 data = random.normal(random.PRNGKey(0), (N, dim)) true_coefs = jnp.arange(1., dim + 1.) logits = jnp.sum(true_coefs * data, axis=-1) labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1)) def model(data, labels): coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim))) logits = jnp.sum(coefs * data, axis=-1) return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels) adam = optim.Adam(0.01) rng_key_init = random.PRNGKey(1) guide = auto_class(model, init_strategy=init_strategy) svi = SVI(model, guide, adam, ELBO()) svi_state = svi.init(rng_key_init, data, labels) def body_fn(i, val): svi_state, loss = svi.update(val, data, labels) return svi_state svi_state = fori_loop(0, 2000, body_fn, svi_state) params = svi.get_params(svi_state) if auto_class not in (AutoIAFNormal, AutoBNAFNormal): median = guide.median(params) assert_allclose(median['coefs'], true_coefs, rtol=0.1) # test .quantile method median = guide.quantiles(params, [0.2, 0.5]) assert_allclose(median['coefs'][1], true_coefs, rtol=0.1) # test .sample_posterior method posterior_samples = guide.sample_posterior(random.PRNGKey(1), params, sample_shape=(1000,)) assert_allclose(jnp.mean(posterior_samples['coefs'], 0), true_coefs, rtol=0.1)
Example #28
Source File: hmc_util.py From numpyro with Apache License 2.0 | 5 votes |
def parametric_draws(subposteriors, num_draws, diagonal=False, rng_key=None): """ Merges subposteriors following (embarrassingly parallel) parametric Monte Carlo algorithm. **References:** 1. *Asymptotically Exact, Embarrassingly Parallel MCMC*, Willie Neiswanger, Chong Wang, Eric Xing :param list subposteriors: a list in which each element is a collection of samples. :param int num_draws: number of draws from the merged posterior. :param bool diagonal: whether to compute weights using variance or covariance, defaults to `False` (using covariance). :param jax.random.PRNGKey rng_key: source of the randomness, defaults to `jax.random.PRNGKey(0)`. :return: a collection of `num_draws` samples with the same data structure as each subposterior. """ rng_key = random.PRNGKey(0) if rng_key is None else rng_key if diagonal: mean, var = parametric(subposteriors, diagonal=True) samples_flat = dist.Normal(mean, jnp.sqrt(var)).sample(rng_key, (num_draws,)) else: mean, cov = parametric(subposteriors, diagonal=False) samples_flat = dist.MultivariateNormal(mean, cov).sample(rng_key, (num_draws,)) _, unravel_fn = ravel_pytree(tree_map(lambda x: x[0], subposteriors[0])) return vmap(lambda x: unravel_fn(x))(samples_flat)
Example #29
Source File: hmc_util.py From numpyro with Apache License 2.0 | 5 votes |
def parametric(subposteriors, diagonal=False): """ Merges subposteriors following (embarrassingly parallel) parametric Monte Carlo algorithm. **References:** 1. *Asymptotically Exact, Embarrassingly Parallel MCMC*, Willie Neiswanger, Chong Wang, Eric Xing :param list subposteriors: a list in which each element is a collection of samples. :param bool diagonal: whether to compute weights using variance or covariance, defaults to `False` (using covariance). :return: the estimated mean and variance/covariance parameters of the joined posterior """ joined_subposteriors = tree_multimap(lambda *args: jnp.stack(args), *subposteriors) joined_subposteriors = vmap(vmap(lambda sample: ravel_pytree(sample)[0]))(joined_subposteriors) submeans = jnp.mean(joined_subposteriors, axis=1) if diagonal: weights = vmap(lambda x: 1 / jnp.var(x, ddof=1, axis=0))(joined_subposteriors) var = 1 / jnp.sum(weights, axis=0) normalized_weights = var * weights # comparing to consensus implementation, we compute weighted mean here mean = jnp.einsum('ij,ij->j', normalized_weights, submeans) return mean, var else: weights = vmap(lambda x: jnp.linalg.inv(jnp.cov(x.T)))(joined_subposteriors) cov = jnp.linalg.inv(jnp.sum(weights, axis=0)) normalized_weights = jnp.matmul(cov, weights) # comparing to consensus implementation, we compute weighted mean here mean = jnp.einsum('ijk,ik->j', normalized_weights, submeans) return mean, cov
Example #30
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 5 votes |
def analyze_dimension(samples, X, Y, dimension, hypers): vmap_args = (samples['msq'], samples['lambda'], samples['eta1'], samples['xisq'], samples['var_obs']) mus, variances = vmap(lambda msq, lam, eta1, xisq, var_obs: compute_singleton_mean_variance(X, Y, dimension, msq, lam, eta1, xisq, hypers['c'], var_obs))(*vmap_args) mean, variance = gaussian_mixture_stats(mus, variances) std = jnp.sqrt(variance) return mean, std # Helper function for analyzing the posterior statistics for coefficient theta_ij