Python jax.numpy.sum() Examples
The following are 30
code examples of jax.numpy.sum().
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_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 #2
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_parameters_from_subsubmodule(): subsublayer = Dense(2) sublayer = Sequential(subsublayer, relu) net = Sequential(sublayer, jnp.sum) inputs = jnp.zeros((1, 3)) params = net.init_parameters(inputs, key=PRNGKey(0)) out = net.apply(params, inputs) subsublayer_params = subsublayer.init_parameters(inputs, key=PRNGKey(0)) params_ = net.parameters_from({subsublayer: subsublayer_params}, inputs) assert_dense_parameters_equal(subsublayer_params, params_[0][0]) out_ = net.apply(params_, inputs) assert out.shape == out_.shape out_ = net.apply_from({subsublayer: subsublayer_params}, inputs) assert out.shape == out_.shape out_ = net.apply_from({subsublayer: subsublayer_params}, inputs, jit=True) assert out.shape == out_.shape
Example #3
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 #4
Source File: test_core.py From jaxnet with Apache License 2.0 | 6 votes |
def test_parameters_from_sharing_between_multiple_parents(): a = Dense(2) b = Sequential(a, jnp.sum) @parametrized def net(inputs): return a(inputs), b(inputs) inputs = jnp.zeros((1, 3)) a_params = a.init_parameters(inputs, key=PRNGKey(0)) out = a.apply(a_params, inputs) params = net.parameters_from({a: a_params}, inputs) assert_dense_parameters_equal(a_params, params.dense) assert_parameters_equal((), params.sequential) assert 2 == len(params) out_, _ = net.apply(params, inputs) assert jnp.array_equal(out, out_)
Example #5
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 #6
Source File: test_distributions.py From numpyro with Apache License 2.0 | 6 votes |
def test_categorical_log_prob_grad(): data = jnp.repeat(jnp.arange(3), 10) def f(x): return dist.Categorical(jax.nn.softmax(x * jnp.arange(1, 4))).log_prob(data).sum() def g(x): return dist.Categorical(logits=x * jnp.arange(1, 4)).log_prob(data).sum() x = 0.5 fx, grad_fx = jax.value_and_grad(f)(x) gx, grad_gx = jax.value_and_grad(g)(x) assert_allclose(fx, gx) assert_allclose(grad_fx, grad_gx, atol=1e-4) ######################################## # Tests for constraints and transforms # ########################################
Example #7
Source File: utils.py From cleverhans with MIT License | 6 votes |
def clip_eta(eta, norm, eps): """ Helper function to clip the perturbation to epsilon norm ball. :param eta: A tensor with the current perturbation. :param norm: Order of the norm (mimics Numpy). Possible values: np.inf or 2. :param eps: Epsilon, bound of the perturbation. """ # Clipping perturbation eta to self.norm norm ball if norm not in [np.inf, 2]: raise ValueError('norm must be np.inf or 2.') axis = list(range(1, len(eta.shape))) avoid_zero_div = 1e-12 if norm == np.inf: eta = np.clip(eta, a_min=-eps, a_max=eps) elif norm == 2: # avoid_zero_div must go inside sqrt to avoid a divide by zero in the gradient through this operation norm = np.sqrt(np.maximum(avoid_zero_div, np.sum(np.square(eta), axis=axis, keepdims=True))) # We must *clip* to within the norm ball, not *normalize* onto the surface of the ball factor = np.minimum(1., np.divide(eps, norm)) eta = eta * factor return eta
Example #8
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 #9
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_diverging(kernel_cls, adapt_step_size): data = random.normal(random.PRNGKey(0), (1000,)) def model(data): loc = numpyro.sample('loc', dist.Normal(0., 1.)) numpyro.sample('obs', dist.Normal(loc, 1), obs=data) kernel = kernel_cls(model, step_size=10., adapt_step_size=adapt_step_size, adapt_mass_matrix=False) num_warmup = num_samples = 1000 mcmc = MCMC(kernel, num_warmup, num_samples) mcmc.warmup(random.PRNGKey(1), data, extra_fields=['diverging'], collect_warmup=True) warmup_divergences = mcmc.get_extra_fields()['diverging'].sum() mcmc.run(random.PRNGKey(2), data, extra_fields=['diverging']) num_divergences = warmup_divergences + mcmc.get_extra_fields()['diverging'].sum() if adapt_step_size: assert num_divergences <= num_warmup else: assert_allclose(num_divergences, num_warmup + num_samples)
Example #10
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 #11
Source File: ppo.py From BERT with Apache License 2.0 | 6 votes |
def masked_entropy(log_probs, mask): """Computes the entropy for the given log-probs. Args: log_probs: (B, T+1, A) log probs mask: (B, T) mask. Returns: Entropy. """ # Cut the last time-step out. lp = log_probs[:, :-1] # Mask out the irrelevant part. lp *= mask[:, :, np.newaxis] # make mask (B, T, 1) p = np.exp(lp) * mask[:, :, np.newaxis] # (B, T, 1) # Average on non-masked part and take negative. return -(np.sum(lp * p) / np.sum(mask))
Example #12
Source File: ppo.py From BERT with Apache License 2.0 | 6 votes |
def approximate_kl(log_prob_new, log_prob_old, mask): """Computes the approximate KL divergence between the old and new log-probs. Args: log_prob_new: (B, T+1, A) log probs new log_prob_old: (B, T+1, A) log probs old mask: (B, T) Returns: Approximate KL. """ diff = log_prob_old - log_prob_new # Cut the last time-step out. diff = diff[:, :-1] # Mask out the irrelevant part. diff *= mask[:, :, np.newaxis] # make mask (B, T, 1) # Average on non-masked part. return np.sum(diff) / np.sum(mask)
Example #13
Source File: test_hmc_util.py From numpyro with Apache License 2.0 | 5 votes |
def kinetic_fn(m_inv, p): return 0.5 * jnp.sum(m_inv * p['x'] ** 2)
Example #14
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def log_prob(self, value): if self._validate_args: self._validate_sample(value) normalize_term = self.total_count * logsumexp(self.logits, axis=-1) \ - gammaln(self.total_count + 1) return jnp.sum(value * self.logits - gammaln(value + 1), axis=-1) - normalize_term
Example #15
Source File: test_optimizers.py From numpyro with Apache License 2.0 | 5 votes |
def loss(params): return jnp.sum(params['x'] ** 2 + params['y'] ** 2)
Example #16
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit): jit_fn = _identity if not jit else jax.jit jax_dist = jax_dist(*params) rng_key = random.PRNGKey(0) samples = jax_dist.sample(key=rng_key, sample_shape=prepend_shape) assert jax_dist.log_prob(samples).shape == prepend_shape + jax_dist.batch_shape if not sp_dist: if isinstance(jax_dist, dist.TruncatedCauchy) or isinstance(jax_dist, dist.TruncatedNormal): low, loc, scale = params high = jnp.inf sp_dist = osp.cauchy if isinstance(jax_dist, dist.TruncatedCauchy) else osp.norm sp_dist = sp_dist(loc, scale) expected = sp_dist.logpdf(samples) - jnp.log(sp_dist.cdf(high) - sp_dist.cdf(low)) assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5) return pytest.skip('no corresponding scipy distn.') if _is_batched_multivariate(jax_dist): pytest.skip('batching not allowed in multivariate distns.') if jax_dist.event_shape and prepend_shape: # >>> d = sp.dirichlet([1.1, 1.1]) # >>> samples = d.rvs(size=(2,)) # >>> d.logpdf(samples) # ValueError: The input vector 'x' must lie within the normal simplex ... pytest.skip('batched samples cannot be scored by multivariate distributions.') sp_dist = sp_dist(*params) try: expected = sp_dist.logpdf(samples) except AttributeError: expected = sp_dist.logpmf(samples) except ValueError as e: # precision issue: jnp.sum(x / jnp.sum(x)) = 0.99999994 != 1 if "The input vector 'x' must lie within the normal simplex." in str(e): samples = samples.copy().astype('float64') samples = samples / samples.sum(axis=-1, keepdims=True) expected = sp_dist.logpdf(samples) else: raise e assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5)
Example #17
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 5 votes |
def test_logistic_regression_x64(kernel_cls): N, dim = 3000, 3 warmup_steps, num_samples = (100000, 100000) if kernel_cls is SA else (1000, 8000) 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 = numpyro.deterministic('logits', jnp.sum(coefs * data, axis=-1)) return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels) if kernel_cls is SA: kernel = SA(model=model, adapt_state_size=9) else: kernel = kernel_cls(model=model, trajectory_length=8, find_heuristic_step_size=True) mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False) mcmc.run(random.PRNGKey(2), labels) mcmc.print_summary() samples = mcmc.get_samples() assert samples['logits'].shape == (num_samples, N) assert_allclose(jnp.mean(samples['coefs'], 0), true_coefs, atol=0.22) if 'JAX_ENABLE_X64' in os.environ: assert samples['coefs'].dtype == jnp.float64
Example #18
Source File: test_example_utils.py From numpyro with Apache License 2.0 | 5 votes |
def test_mnist_data_load(): def mean_pixels(i, mean_pix): batch, _ = fetch(i, idx) return mean_pix + jnp.sum(batch) / batch.size init, fetch = load_dataset(MNIST, batch_size=128, split='train') num_batches, idx = init() assert fori_loop(0, num_batches, mean_pixels, jnp.float32(0.)) / num_batches < 0.15
Example #19
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def log_prob(self, value): init_prob = Normal(0., self.scale).log_prob(value[..., 0]) scale = jnp.expand_dims(self.scale, -1) step_probs = Normal(value[..., :-1], scale).log_prob(value[..., 1:]) return init_prob + jnp.sum(step_probs, axis=-1)
Example #20
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def variance(self): raw_variance = jnp.square(self.cov_factor).sum(-1) + self.cov_diag return jnp.broadcast_to(raw_variance, self.batch_shape + self.event_shape)
Example #21
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def variance(self): con0 = jnp.sum(self.concentration, axis=-1, keepdims=True) return self.concentration * (con0 - self.concentration) / (con0 ** 2 * (con0 + 1))
Example #22
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def mean(self): return self.concentration / jnp.sum(self.concentration, axis=-1, keepdims=True)
Example #23
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def log_prob(self, value): normalize_term = (jnp.sum(gammaln(self.concentration), axis=-1) - gammaln(jnp.sum(self.concentration, axis=-1))) return jnp.sum(jnp.log(value) * (self.concentration - 1.), axis=-1) - normalize_term
Example #24
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def sample(self, key, sample_shape=()): shape = sample_shape + self.batch_shape + self.event_shape gamma_samples = random.gamma(key, self.concentration, shape=shape) samples = gamma_samples / jnp.sum(gamma_samples, axis=-1, keepdims=True) return jnp.clip(samples, a_min=jnp.finfo(samples).tiny, a_max=1 - jnp.finfo(samples).eps)
Example #25
Source File: util.py From numpyro with Apache License 2.0 | 5 votes |
def _categorical(key, p, shape): # this implementation is fast when event shape is small, and slow otherwise # Ref: https://stackoverflow.com/a/34190035 shape = shape or p.shape[:-1] s = jnp.cumsum(p, axis=-1) r = random.uniform(key, shape=shape + (1,)) # FIXME: replace this computation by using binary search as suggested in the above # reference. A while_loop + vmap for a reshaped 2D array would be enough. return jnp.sum(s < r, axis=-1)
Example #26
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def log_abs_det_jacobian(self, x, y, intermediates=None): # Ref: https://mc-stan.org/docs/2_19/reference-manual/simplex-transform-section.html # |det|(J) = Product(y * (1 - z)) x = x - jnp.log(x.shape[-1] - jnp.arange(x.shape[-1])) z = jnp.clip(expit(x), a_min=jnp.finfo(x.dtype).tiny) # XXX we use the identity 1 - z = z * exp(-x) to not worry about # the case z ~ 1 return jnp.sum(jnp.log(y[..., :-1] * z) - x, axis=-1)
Example #27
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def log_abs_det_jacobian(self, x, y, intermediates=None): return jnp.sum(x[..., 1:], -1)
Example #28
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def log_abs_det_jacobian(self, x, y, intermediates=None): # the jacobian is diagonal, so logdet is the sum of diagonal `exp` transform n = round((math.sqrt(1 + 8 * x.shape[-1]) - 1) / 2) return x[..., -n:].sum(-1)
Example #29
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def log_abs_det_jacobian(self, x, y, intermediates=None): return jnp.broadcast_to(jnp.log(jnp.diagonal(self.scale_tril, axis1=-2, axis2=-1)).sum(-1), jnp.shape(x)[:-1])
Example #30
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def log_abs_det_jacobian(self, x, y, intermediates=None): # NB: because domain and codomain are two spaces with different dimensions, determinant of # Jacobian is not well-defined. Here we return `log_abs_det_jacobian` of `x` and the # flatten lower triangular part of `y`. # stick_breaking_logdet = log(y / r) = log(z_cumprod) (modulo right shifted) z1m_cumprod = 1 - jnp.cumsum(y * y, axis=-1) # by taking diagonal=-2, we don't need to shift z_cumprod to the right # NB: diagonal=-2 works fine for (2 x 2) matrix, where we get an empty array z1m_cumprod_tril = matrix_to_tril_vec(z1m_cumprod, diagonal=-2) stick_breaking_logdet = 0.5 * jnp.sum(jnp.log(z1m_cumprod_tril), axis=-1) tanh_logdet = -2 * jnp.sum(x + softplus(-2 * x) - jnp.log(2.), axis=-1) return stick_breaking_logdet + tanh_logdet