Python jax.vmap() Examples
The following are 24
code examples of jax.vmap().
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
, or try the search function
.
![](https://www.programcreek.com/common/static/images/search.png)
Example #1
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def log_likelihood(model, posterior_samples, *args, **kwargs): """ (EXPERIMENTAL INTERFACE) Returns log likelihood at observation nodes of model, given samples of all latent variables. :param model: Python callable containing Pyro primitives. :param dict posterior_samples: dictionary of samples from the posterior. :param args: model arguments. :param kwargs: model kwargs. :return: dict of log likelihoods at observation sites. """ def single_loglik(samples): model_trace = trace(substitute(model, samples)).get_trace(*args, **kwargs) return {name: site['fn'].log_prob(site['value']) for name, site in model_trace.items() if site['type'] == 'sample' and site['is_observed']} return vmap(single_loglik)(posterior_samples)
Example #2
Source File: test_distributions_util.py From numpyro with Apache License 2.0 | 6 votes |
def test_binop_batch_rule(prim): bx = jnp.array([1., 2., 3.]) by = jnp.array([2., 3., 4.]) x = jnp.array(1.) y = jnp.array(2.) actual_bx_by = vmap(lambda x, y: prim(x, y))(bx, by) for i in range(3): assert_allclose(actual_bx_by[i], prim(bx[i], by[i])) actual_x_by = vmap(lambda y: prim(x, y))(by) for i in range(3): assert_allclose(actual_x_by[i], prim(x, by[i])) actual_bx_y = vmap(lambda x: prim(x, y))(bx) for i in range(3): assert_allclose(actual_bx_y[i], prim(bx[i], y))
Example #3
Source File: test_nn.py From numpyro with Apache License 2.0 | 6 votes |
def test_block_neural_arn(input_dim, hidden_factors, residual, batch_shape): arn_init, arn = BlockNeuralAutoregressiveNN(input_dim, hidden_factors, residual) rng = random.PRNGKey(0) input_shape = batch_shape + (input_dim,) out_shape, init_params = arn_init(rng, input_shape) assert out_shape == input_shape x = random.normal(random.PRNGKey(1), input_shape) output, logdet = arn(init_params, x) assert output.shape == input_shape assert logdet.shape == input_shape if len(batch_shape) == 1: jac = vmap(jacfwd(lambda x: arn(init_params, x)[0]))(x) else: jac = jacfwd(lambda x: arn(init_params, x)[0])(x) assert_allclose(logdet.sum(-1), jnp.linalg.slogdet(jac)[1], rtol=1e-6) # make sure jacobians are lower triangular assert np.sum(np.abs(np.triu(jac, k=1))) == 0.0 assert np.all(np.abs(matrix_to_tril_vec(jac)) > 0)
Example #4
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def _multinomial(key, p, n, n_max, shape=()): if jnp.shape(n) != jnp.shape(p)[:-1]: broadcast_shape = lax.broadcast_shapes(jnp.shape(n), jnp.shape(p)[:-1]) n = jnp.broadcast_to(n, broadcast_shape) p = jnp.broadcast_to(p, broadcast_shape + jnp.shape(p)[-1:]) shape = shape or p.shape[:-1] # get indices from categorical distribution then gather the result indices = categorical(key, p, (n_max,) + shape) # mask out values when counts is heterogeneous if jnp.ndim(n) > 0: mask = promote_shapes(jnp.arange(n_max) < jnp.expand_dims(n, -1), shape=shape + (n_max,))[0] mask = jnp.moveaxis(mask, -1, 0).astype(indices.dtype) excess = jnp.concatenate([jnp.expand_dims(n_max - n, -1), jnp.zeros(jnp.shape(n) + (p.shape[-1] - 1,))], -1) else: mask = 1 excess = 0 # NB: we transpose to move batch shape to the front indices_2D = (jnp.reshape(indices * mask, (n_max, -1,))).T samples_2D = vmap(_scatter_add_one, (0, 0, 0))(jnp.zeros((indices_2D.shape[0], p.shape[-1]), dtype=indices.dtype), jnp.expand_dims(indices_2D, axis=-1), jnp.ones(indices_2D.shape, dtype=indices.dtype)) return jnp.reshape(samples_2D, shape + p.shape[-1:]) - excess
Example #5
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def _predictive(rng_key, model, posterior_samples, num_samples, return_sites=None, parallel=True, model_args=(), model_kwargs={}): rng_keys = random.split(rng_key, num_samples) def single_prediction(val): rng_key, samples = val model_trace = trace(seed(substitute(model, samples), rng_key)).get_trace( *model_args, **model_kwargs) if return_sites is not None: if return_sites == '': sites = {k for k, site in model_trace.items() if site['type'] != 'plate'} else: sites = return_sites else: sites = {k for k, site in model_trace.items() if (site['type'] == 'sample' and k not in samples) or (site['type'] == 'deterministic')} return {name: site['value'] for name, site in model_trace.items() if name in sites} if parallel: return vmap(single_prediction)((rng_keys, posterior_samples)) else: return lax.map(single_prediction, (rng_keys, posterior_samples))
Example #6
Source File: test_modules.py From jaxnet with Apache License 2.0 | 6 votes |
def test_Batched(): out_dim = 1 @parametrized def unbatched_dense(input): kernel = parameter((out_dim, input.shape[-1]), ones) bias = parameter((out_dim,), ones) return jnp.dot(kernel, input) + bias batch_size = 4 unbatched_params = unbatched_dense.init_parameters(jnp.zeros(2), key=PRNGKey(0)) out = unbatched_dense.apply(unbatched_params, jnp.ones(2)) assert jnp.array([3.]) == out dense_apply = vmap(unbatched_dense.apply, (None, 0)) out_batched_ = dense_apply(unbatched_params, jnp.ones((batch_size, 2))) assert jnp.array_equal(jnp.stack([out] * batch_size), out_batched_) dense = Batched(unbatched_dense) params = dense.init_parameters(jnp.ones((batch_size, 2)), key=PRNGKey(0)) assert_parameters_equal((unbatched_params,), params) out_batched = dense.apply(params, jnp.ones((batch_size, 2))) assert jnp.array_equal(out_batched_, out_batched)
Example #7
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 #8
Source File: modules.py From jaxnet with Apache License 2.0 | 5 votes |
def Batched(unbatched_model: parametrized, batch_dim=0): @parametrized def batched(*batched_args): args = tree_map(lambda x: x[0], batched_args) params = Parameter(lambda key: unbatched_model.init_parameters(*args, key=key), 'model')() batched_apply = vmap(partial(unbatched_model.apply, params), batch_dim) return batched_apply(*batched_args) return batched
Example #9
Source File: test_handlers.py From numpyro with Apache License 2.0 | 5 votes |
def test_seed(): def _sample(): x = numpyro.sample('x', dist.Normal(0., 1.)) y = numpyro.sample('y', dist.Normal(1., 2.)) return jnp.stack([x, y]) xs = [] for i in range(100): with handlers.seed(rng_seed=i): xs.append(_sample()) xs = jnp.stack(xs) ys = vmap(lambda rng_key: handlers.seed(lambda: _sample(), rng_key)())(jnp.arange(100)) assert_allclose(xs, ys, atol=1e-6)
Example #10
Source File: bnn.py From numpyro with Apache License 2.0 | 5 votes |
def main(args): N, D_X, D_H = args.num_data, 3, args.num_hidden X, Y, X_test = get_data(N=N, D_X=D_X) # do inference rng_key, rng_key_predict = random.split(random.PRNGKey(0)) samples = run_inference(model, args, rng_key, X, Y, D_H) # predict Y_test at inputs X_test vmap_args = (samples, random.split(rng_key_predict, args.num_samples * args.num_chains)) predictions = vmap(lambda samples, rng_key: predict(model, rng_key, samples, X_test, D_H))(*vmap_args) predictions = predictions[..., 0] # compute mean prediction and confidence interval around median mean_prediction = jnp.mean(predictions, 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[:, 1], Y[:, 0], 'kx') # plot 90% confidence level of predictions ax.fill_between(X_test[:, 1], percentiles[0, :], percentiles[1, :], color='lightblue') # plot mean prediction ax.plot(X_test[:, 1], mean_prediction, 'blue', ls='solid', lw=2.0) ax.set(xlabel="X", ylabel="Y", title="Mean predictions with 90% CI") plt.savefig('bnn_plot.pdf') plt.tight_layout()
Example #11
Source File: util.py From numpyro with Apache License 2.0 | 5 votes |
def _binomial(key, p, n, shape): shape = shape or lax.broadcast_shapes(jnp.shape(p), jnp.shape(n)) # reshape to map over axis 0 p = jnp.reshape(jnp.broadcast_to(p, shape), -1) n = jnp.reshape(jnp.broadcast_to(n, shape), -1) key = random.split(key, jnp.size(p)) if xla_bridge.get_backend().platform == 'cpu': ret = lax.map(lambda x: _binomial_dispatch(*x), (key, p, n)) else: ret = vmap(lambda *x: _binomial_dispatch(*x))(key, p, n) return jnp.reshape(ret, shape)
Example #12
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def __call__(self, x): batch_shape = x.shape[:-1] if batch_shape: unpacked = vmap(self.unpack_fn)(x.reshape((-1,) + x.shape[-1:])) return tree_map(lambda z: jnp.reshape(z, batch_shape + z.shape[1:]), unpacked) else: return self.unpack_fn(x)
Example #13
Source File: elbo.py From numpyro with Apache License 2.0 | 5 votes |
def loss(self, rng_key, param_map, model, guide, *args, **kwargs): """ Evaluates the Renyi ELBO with an estimator that uses num_particles many samples/particles. :param jax.random.PRNGKey rng_key: random number generator seed. :param dict param_map: dictionary of current parameter values keyed by site name. :param model: Python callable with NumPyro primitives for the model. :param guide: Python callable with NumPyro primitives for the guide. :param args: arguments to the model / guide (these can possibly vary during the course of fitting). :param kwargs: keyword arguments to the model / guide (these can possibly vary during the course of fitting). :returns: negative of the Renyi Evidence Lower Bound (ELBO) to be minimized. """ def single_particle_elbo(rng_key): model_seed, guide_seed = random.split(rng_key) seeded_model = seed(model, model_seed) seeded_guide = seed(guide, guide_seed) guide_log_density, guide_trace = log_density(seeded_guide, args, kwargs, param_map) # NB: we only want to substitute params not available in guide_trace model_param_map = {k: v for k, v in param_map.items() if k not in guide_trace} seeded_model = replay(seeded_model, guide_trace) model_log_density, _ = log_density(seeded_model, args, kwargs, model_param_map) # log p(z) - log q(z) elbo = model_log_density - guide_log_density return elbo rng_keys = random.split(rng_key, self.num_particles) elbos = vmap(single_particle_elbo)(rng_keys) scaled_elbos = (1. - self.alpha) * elbos avg_log_exp = logsumexp(scaled_elbos) - jnp.log(self.num_particles) weights = jnp.exp(scaled_elbos - avg_log_exp) renyi_elbo = avg_log_exp / (1. - self.alpha) weighted_elbo = jnp.dot(stop_gradient(weights), elbos) / self.num_particles return - (stop_gradient(renyi_elbo - weighted_elbo) + weighted_elbo)
Example #14
Source File: elbo.py From numpyro with Apache License 2.0 | 5 votes |
def loss(self, rng_key, param_map, model, guide, *args, **kwargs): """ Evaluates the ELBO with an estimator that uses num_particles many samples/particles. :param jax.random.PRNGKey rng_key: random number generator seed. :param dict param_map: dictionary of current parameter values keyed by site name. :param model: Python callable with NumPyro primitives for the model. :param guide: Python callable with NumPyro primitives for the guide. :param args: arguments to the model / guide (these can possibly vary during the course of fitting). :param kwargs: keyword arguments to the model / guide (these can possibly vary during the course of fitting). :return: negative of the Evidence Lower Bound (ELBO) to be minimized. """ def single_particle_elbo(rng_key): model_seed, guide_seed = random.split(rng_key) seeded_model = seed(model, model_seed) seeded_guide = seed(guide, guide_seed) guide_log_density, guide_trace = log_density(seeded_guide, args, kwargs, param_map) seeded_model = replay(seeded_model, guide_trace) model_log_density, _ = log_density(seeded_model, args, kwargs, param_map) # log p(z) - log q(z) elbo = model_log_density - guide_log_density return elbo # Return (-elbo) since by convention we do gradient descent on a loss and # the ELBO is a lower bound that needs to be maximized. if self.num_particles == 1: return - single_particle_elbo(rng_key) else: rng_keys = random.split(rng_key, self.num_particles) return - jnp.mean(vmap(single_particle_elbo)(rng_keys))
Example #15
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 5 votes |
def analyze_pair_of_dimensions(samples, X, Y, dim1, dim2, 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_pairwise_mean_variance(X, Y, dim1, dim2, msq, lam, eta1, xisq, hypers['c'], var_obs))(*vmap_args) mean, variance = gaussian_mixture_stats(mus, variances) std = jnp.sqrt(variance) return mean, std
Example #16
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def test_log_prob_LKJCholesky(dimension, concentration): # We will test against the fact that LKJCorrCholesky can be seen as a # TransformedDistribution with base distribution is a distribution of partial # correlations in C-vine method (modulo an affine transform to change domain from (0, 1) # to (1, 0)) and transform is a signed stick-breaking process. d = dist.LKJCholesky(dimension, concentration, sample_method="cvine") beta_sample = d._beta.sample(random.PRNGKey(0)) beta_log_prob = jnp.sum(d._beta.log_prob(beta_sample)) partial_correlation = 2 * beta_sample - 1 affine_logdet = beta_sample.shape[-1] * jnp.log(2) sample = signed_stick_breaking_tril(partial_correlation) # compute signed stick breaking logdet inv_tanh = lambda t: jnp.log((1 + t) / (1 - t)) / 2 # noqa: E731 inv_tanh_logdet = jnp.sum(jnp.log(vmap(grad(inv_tanh))(partial_correlation))) unconstrained = inv_tanh(partial_correlation) corr_cholesky_logdet = biject_to(constraints.corr_cholesky).log_abs_det_jacobian( unconstrained, sample, ) signed_stick_breaking_logdet = corr_cholesky_logdet + inv_tanh_logdet actual_log_prob = d.log_prob(sample) expected_log_prob = beta_log_prob - affine_logdet - signed_stick_breaking_logdet assert_allclose(actual_log_prob, expected_log_prob, rtol=2e-5) assert_allclose(jax.jit(d.log_prob)(sample), d.log_prob(sample), atol=1e-7)
Example #17
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
Example #18
Source File: mcmc.py From numpyro with Apache License 2.0 | 5 votes |
def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={}): # non-vectorized if rng_key.ndim == 1: rng_key, rng_key_init_model = random.split(rng_key) # vectorized else: rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1) # we need only a single key for initializing PE / constraints fn rng_key_init_model = rng_key_init_model[0] init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params) if self._potential_fn and init_params is None: raise ValueError('Valid value of `init_params` must be provided with' ' `potential_fn`.') # NB: init args is different from HMC sa_init_fn = lambda init_params, rng_key: self._init_fn( # noqa: E731 init_params, num_warmup=num_warmup, adapt_state_size=self._adapt_state_size, dense_mass=self._dense_mass, rng_key=rng_key, model_args=model_args, model_kwargs=model_kwargs, ) if rng_key.ndim == 1: init_state = sa_init_fn(init_params, rng_key) else: init_state = vmap(sa_init_fn)(init_params, rng_key) sample_fn = vmap(self._sample_fn, in_axes=(0, None, None)) self._sample_fn = sample_fn return init_state
Example #19
Source File: mcmc.py From numpyro with Apache License 2.0 | 5 votes |
def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={}): # non-vectorized if rng_key.ndim == 1: rng_key, rng_key_init_model = random.split(rng_key) # vectorized else: rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1) init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params) if self._potential_fn and init_params is None: raise ValueError('Valid value of `init_params` must be provided with' ' `potential_fn`.') hmc_init_fn = lambda init_params, rng_key: self._init_fn( # noqa: E731 init_params, num_warmup=num_warmup, step_size=self._step_size, adapt_step_size=self._adapt_step_size, adapt_mass_matrix=self._adapt_mass_matrix, dense_mass=self._dense_mass, target_accept_prob=self._target_accept_prob, trajectory_length=self._trajectory_length, max_tree_depth=self._max_tree_depth, find_heuristic_step_size=self._find_heuristic_step_size, model_args=model_args, model_kwargs=model_kwargs, rng_key=rng_key, ) if rng_key.ndim == 1: init_state = hmc_init_fn(init_params, rng_key) else: # XXX it is safe to run hmc_init_fn under vmap despite that hmc_init_fn changes some # nonlocal variables: momentum_generator, wa_update, trajectory_len, max_treedepth, # wa_steps because those variables do not depend on traced args: init_params, rng_key. init_state = vmap(hmc_init_fn)(init_params, rng_key) sample_fn = vmap(self._sample_fn, in_axes=(0, None, None)) self._sample_fn = sample_fn return init_state
Example #20
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 #21
Source File: tabular_irl.py From imitation with MIT License | 5 votes |
def grads(self, inputs): in_grad_partial = jax.partial(self._net_grads, self._net_params) grad_vmap = jax.vmap(in_grad_partial) rich_grads = grad_vmap(inputs) flat_grads = np.asarray(self._flatten_batch(rich_grads)) assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0] return flat_grads
Example #22
Source File: fast_gradient_method.py From cleverhans with MIT License | 4 votes |
def fast_gradient_method(model_fn, x, eps, norm, clip_min=None, clip_max=None, y=None, targeted=False): """ JAX implementation of the Fast Gradient Method. :param model_fn: a callable that takes an input tensor and returns the model logits. :param x: input tensor. :param eps: epsilon (input variation parameter); see https://arxiv.org/abs/1412.6572. :param norm: Order of the norm (mimics NumPy). Possible values: np.inf or 2. :param clip_min: (optional) float. Minimum float value for adversarial example components. :param clip_max: (optional) float. Maximum float value for adversarial example components. :param y: (optional) Tensor with one-hot true labels. If targeted is true, then provide the target one-hot label. Otherwise, only provide this parameter if you'd like to use true labels when crafting adversarial samples. Otherwise, model predictions are used as labels to avoid the "label leaking" effect (explained in this paper: https://arxiv.org/abs/1611.01236). Default is None. This argument does not have to be a binary one-hot label (e.g., [0, 1, 0, 0]), it can be floating points values that sum up to 1 (e.g., [0.05, 0.85, 0.05, 0.05]). :param targeted: (optional) bool. Is the attack targeted or untargeted? Untargeted, the default, will try to make the label incorrect. Targeted will instead try to move in the direction of being more like y. :return: a tensor for the adversarial example """ if norm not in [np.inf, 2]: raise ValueError("Norm order must be either np.inf or 2.") if y is None: # Using model predictions as ground truth to avoid label leaking x_labels = np.argmax(model_fn(x), 1) y = one_hot(x_labels, 10) def loss_adv(image, label): pred = model_fn(image[None]) loss = - np.sum(logsoftmax(pred) * label) if targeted: loss = -loss return loss grads_fn = vmap(grad(loss_adv), in_axes=(0, 0), out_axes=0) grads = grads_fn(x, y) axis = list(range(1, len(grads.shape))) avoid_zero_div = 1e-12 if norm == np.inf: perturbation = eps * np.sign(grads) elif norm == 1: raise NotImplementedError("L_1 norm has not been implemented yet.") elif norm == 2: square = np.maximum(avoid_zero_div, np.sum(np.square(grads), axis=axis, keepdims=True)) perturbation = grads / np.sqrt(square) adv_x = x + perturbation # If clipping is needed, reset all values outside of [clip_min, clip_max] if (clip_min is not None) or (clip_max is not None): # We don't currently support one-sided clipping assert clip_min is not None and clip_max is not None adv_x = np.clip(adv_x, a_min=clip_min, a_max=clip_max) return adv_x
Example #23
Source File: hmc_util.py From numpyro with Apache License 2.0 | 4 votes |
def consensus(subposteriors, num_draws=None, diagonal=False, rng_key=None): """ Merges subposteriors following consensus Monte Carlo algorithm. **References:** 1. *Bayes and big data: The consensus Monte Carlo algorithm*, Steven L. Scott, Alexander W. Blocker, Fernando V. Bonassi, Hugh A. Chipman, Edward I. George, Robert E. McCulloch :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: if `num_draws` is None, merges subposteriors without resampling; otherwise, returns a collection of `num_draws` samples with the same data structure as each subposterior. """ # stack subposteriors joined_subposteriors = tree_multimap(lambda *args: jnp.stack(args), *subposteriors) # shape of joined_subposteriors: n_subs x n_samples x sample_shape joined_subposteriors = vmap(vmap(lambda sample: ravel_pytree(sample)[0]))(joined_subposteriors) if num_draws is not None: rng_key = random.PRNGKey(0) if rng_key is None else rng_key # randomly gets num_draws from subposteriors n_subs = len(subposteriors) n_samples = tree_flatten(subposteriors[0])[0][0].shape[0] # shape of draw_idxs: n_subs x num_draws x sample_shape draw_idxs = random.randint(rng_key, shape=(n_subs, num_draws), minval=0, maxval=n_samples) joined_subposteriors = vmap(lambda x, idx: x[idx])(joined_subposteriors, draw_idxs) if diagonal: # compute weights for each subposterior (ref: Section 3.1 of [1]) weights = vmap(lambda x: 1 / jnp.var(x, ddof=1, axis=0))(joined_subposteriors) normalized_weights = weights / jnp.sum(weights, axis=0) # get weighted samples samples_flat = jnp.einsum('ij,ikj->kj', normalized_weights, joined_subposteriors) else: weights = vmap(lambda x: jnp.linalg.inv(jnp.cov(x.T)))(joined_subposteriors) normalized_weights = jnp.matmul(jnp.linalg.inv(jnp.sum(weights, axis=0)), weights) samples_flat = jnp.einsum('ijk,ilk->lj', normalized_weights, joined_subposteriors) # unravel_fn acts on 1 sample of a subposterior _, unravel_fn = ravel_pytree(tree_map(lambda x: x[0], subposteriors[0])) return vmap(lambda x: unravel_fn(x))(samples_flat)
Example #24
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 4 votes |
def main(args): X, Y, expected_thetas, expected_pairwise = get_data(N=args.num_data, P=args.num_dimensions, S=args.active_dimensions) # setup hyperparameters hypers = {'expected_sparsity': max(1.0, args.num_dimensions / 10), 'alpha1': 3.0, 'beta1': 1.0, 'alpha2': 3.0, 'beta2': 1.0, 'alpha3': 1.0, 'c': 1.0, 'alpha_obs': 3.0, 'beta_obs': 1.0} # do inference rng_key = random.PRNGKey(0) samples = run_inference(model, args, rng_key, X, Y, hypers) # compute the mean and square root variance of each coefficient theta_i means, stds = vmap(lambda dim: analyze_dimension(samples, X, Y, dim, hypers))(jnp.arange(args.num_dimensions)) print("Coefficients theta_1 to theta_%d used to generate the data:" % args.active_dimensions, expected_thetas) print("The single quadratic coefficient theta_{1,2} used to generate the data:", expected_pairwise) active_dimensions = [] for dim, (mean, std) in enumerate(zip(means, stds)): # we mark the dimension as inactive if the interval [mean - 3 * std, mean + 3 * std] contains zero lower, upper = mean - 3.0 * std, mean + 3.0 * std inactive = "inactive" if lower < 0.0 and upper > 0.0 else "active" if inactive == "active": active_dimensions.append(dim) print("[dimension %02d/%02d] %s:\t%.2e +- %.2e" % (dim + 1, args.num_dimensions, inactive, mean, std)) print("Identified a total of %d active dimensions; expected %d." % (len(active_dimensions), args.active_dimensions)) # Compute the mean and square root variance of coefficients theta_ij for i,j active dimensions. # Note that the resulting numbers are only meaningful for i != j. if len(active_dimensions) > 0: dim_pairs = jnp.array(list(itertools.product(active_dimensions, active_dimensions))) means, stds = vmap(lambda dim_pair: analyze_pair_of_dimensions(samples, X, Y, dim_pair[0], dim_pair[1], hypers))(dim_pairs) for dim_pair, mean, std in zip(dim_pairs, means, stds): dim1, dim2 = dim_pair if dim1 >= dim2: continue lower, upper = mean - 3.0 * std, mean + 3.0 * std if not (lower < 0.0 and upper > 0.0): format_str = "Identified pairwise interaction between dimensions %d and %d: %.2e +- %.2e" print(format_str % (dim1 + 1, dim2 + 1, mean, std)) # Draw a single sample of coefficients theta from the posterior, where we return all singleton # coefficients theta_i and pairwise coefficients theta_ij for i, j active dimensions. We use the # final MCMC sample obtained from the HMC sampler. thetas = sample_theta_space(X, Y, active_dimensions, samples['msq'][-1], samples['lambda'][-1], samples['eta1'][-1], samples['xisq'][-1], hypers['c'], samples['var_obs'][-1]) print("Single posterior sample theta:\n", thetas)