Python jax.numpy.arange() Examples
The following are 30
code examples of jax.numpy.arange().
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: tke_jax.py From pyhpc-benchmarks with The Unlicense | 6 votes |
def solve_implicit(ks, a, b, c, d, b_edge=None, d_edge=None): land_mask = (ks >= 0)[:, :, np.newaxis] edge_mask = land_mask & (np.arange(a.shape[2])[np.newaxis, np.newaxis, :] == ks[:, :, np.newaxis]) water_mask = land_mask & (np.arange(a.shape[2])[np.newaxis, np.newaxis, :] >= ks[:, :, np.newaxis]) a_tri = water_mask * a * np.logical_not(edge_mask) b_tri = where(water_mask, b, 1.) if b_edge is not None: b_tri = where(edge_mask, b_edge, b_tri) c_tri = water_mask * c d_tri = water_mask * d if d_edge is not None: d_tri = where(edge_mask, d_edge, d_tri) return solve_tridiag(a_tri, b_tri, c_tri, d_tri), water_mask
Example #2
Source File: ppo.py From BERT with Apache License 2.0 | 6 votes |
def chosen_probabs(probab_observations, actions): """Picks out the probabilities of the actions along batch and time-steps. Args: probab_observations: ndarray of shape `[B, T+1, A]`, where probab_observations[b, t, i] contains the log-probability of action = i at the t^th time-step in the b^th trajectory. actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which action was chosen in the b^th trajectory's t^th time-step. Returns: `[B, T]` ndarray with the log-probabilities of the chosen actions. """ B, T = actions.shape # pylint: disable=invalid-name assert (B, T + 1) == probab_observations.shape[:2] return probab_observations[np.arange(B)[:, None], np.arange(T), actions]
Example #3
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 #4
Source File: ode.py From numpyro with Apache License 2.0 | 6 votes |
def model(N, y=None): """ :param int N: number of measurement times :param numpy.ndarray y: measured populations with shape (N, 2) """ # initial population z_init = numpyro.sample("z_init", dist.LogNormal(jnp.log(10), 1), sample_shape=(2,)) # measurement times ts = jnp.arange(float(N)) # parameters alpha, beta, gamma, delta of dz_dt theta = numpyro.sample( "theta", dist.TruncatedNormal(low=0., loc=jnp.array([0.5, 0.05, 1.5, 0.05]), scale=jnp.array([0.5, 0.05, 0.5, 0.05]))) # integrate dz/dt, the result will have shape N x 2 z = odeint(dz_dt, z_init, ts, theta, rtol=1e-5, atol=1e-3, mxstep=500) # measurement errors, we expect that measured hare has larger error than measured lynx sigma = numpyro.sample("sigma", dist.Exponential(jnp.array([1, 2]))) # measured populations (in log scale) numpyro.sample("y", dist.Normal(jnp.log(z), sigma), obs=y)
Example #5
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 #6
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 #7
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 #8
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 #9
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): total_count = jnp.amax(self.total_count) if not_jax_tracer(total_count): # NB: the error can't be raised if inhomogeneous issue happens when tracing if jnp.amin(self.total_count) != total_count: raise NotImplementedError("Inhomogeneous total count not supported" " by `enumerate_support`.") values = jnp.arange(total_count + 1).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #10
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def __init__(self, dimension, concentration=1., sample_method='onion', validate_args=None): if dimension < 2: raise ValueError("Dimension must be greater than or equal to 2.") self.dimension = dimension self.concentration = concentration batch_shape = jnp.shape(concentration) event_shape = (dimension, dimension) # We construct base distributions to generate samples for each method. # The purpose of this base distribution is to generate a distribution for # correlation matrices which is propotional to `det(M)^{\eta - 1}`. # (note that this is not a unique way to define base distribution) # Both of the following methods have marginal distribution of each off-diagonal # element of sampled correlation matrices is Beta(eta + (D-2) / 2, eta + (D-2) / 2) # (up to a linear transform: x -> 2x - 1) Dm1 = self.dimension - 1 marginal_concentration = concentration + 0.5 * (self.dimension - 2) offset = 0.5 * jnp.arange(Dm1) if sample_method == 'onion': # The following construction follows from the algorithm in Section 3.2 of [1]: # NB: in [1], the method for case k > 1 can also work for the case k = 1. beta_concentration0 = jnp.expand_dims(marginal_concentration, axis=-1) - offset beta_concentration1 = offset + 0.5 self._beta = Beta(beta_concentration1, beta_concentration0) elif sample_method == 'cvine': # The following construction follows from the algorithm in Section 2.4 of [1]: # offset_tril is [0, 1, 1, 2, 2, 2,...] / 2 offset_tril = matrix_to_tril_vec(jnp.broadcast_to(offset, (Dm1, Dm1))) beta_concentration = jnp.expand_dims(marginal_concentration, axis=-1) - offset_tril self._beta = Beta(beta_concentration, beta_concentration) else: raise ValueError("`method` should be one of 'cvine' or 'onion'.") self.sample_method = sample_method super(LKJCholesky, self).__init__(batch_shape=batch_shape, event_shape=event_shape, validate_args=validate_args)
Example #11
Source File: util.py From numpyro with Apache License 2.0 | 5 votes |
def vec_to_tril_matrix(t, diagonal=0): # NB: the following formula only works for diagonal <= 0 n = round((math.sqrt(1 + 8 * t.shape[-1]) - 1) / 2) - diagonal n2 = n * n idx = jnp.reshape(jnp.arange(n2), (n, n))[jnp.tril_indices(n, diagonal)] x = lax.scatter_add(jnp.zeros(t.shape[:-1] + (n2,)), jnp.expand_dims(idx, axis=-1), t, lax.ScatterDimensionNumbers(update_window_dims=range(t.ndim - 1), inserted_window_dims=(t.ndim - 1,), scatter_dims_to_operand_dims=(t.ndim - 1,))) return jnp.reshape(x, x.shape[:-1] + (n, n))
Example #12
Source File: utils.py From cleverhans with MIT License | 5 votes |
def one_hot(x, k, dtype=np.float32): """Create a one-hot encoding of x of size k.""" return np.array(x[:, None] == np.arange(k), dtype)
Example #13
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def sample(self, key, sample_shape=()): denom = jnp.square(jnp.arange(0.5, self.num_gamma_variates)) x = random.gamma(key, jnp.ones(self.batch_shape + sample_shape + (self.num_gamma_variates,))) x = jnp.sum(x / denom, axis=-1) return jnp.clip(x * (0.5 / jnp.pi ** 2), a_max=self.truncation_point)
Example #14
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def log_prob(self, value): value = value[..., None] all_indices = jnp.arange(0, self.num_log_prob_terms) two_n_plus_one = 2.0 * all_indices + 1.0 log_terms = jnp.log(two_n_plus_one) - 1.5 * jnp.log(value) - 0.125 * jnp.square(two_n_plus_one) / value even_terms = jnp.take(log_terms, all_indices[::2], axis=-1) odd_terms = jnp.take(log_terms, all_indices[1::2], axis=-1) sum_even = jnp.exp(logsumexp(even_terms, axis=-1)) sum_odd = jnp.exp(logsumexp(odd_terms, axis=-1)) return jnp.log(sum_even - sum_odd) - 0.5 * jnp.log(2.0 * jnp.pi)
Example #15
Source File: conjugate.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): total_count = jnp.amax(self.total_count) if not_jax_tracer(total_count): # NB: the error can't be raised if inhomogeneous issue happens when tracing if jnp.amin(self.total_count) != total_count: raise NotImplementedError("Inhomogeneous total count not supported" " by `enumerate_support`.") values = jnp.arange(total_count + 1).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #16
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): values = jnp.arange(2).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #17
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): values = jnp.arange(2).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #18
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 #19
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): values = jnp.arange(self.probs.shape[-1]).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #20
Source File: discrete.py From numpyro with Apache License 2.0 | 5 votes |
def enumerate_support(self, expand=True): values = jnp.arange(self.logits.shape[-1]).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
Example #21
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def test_gamma_poisson_log_prob(shape): gamma_conc = np.exp(np.random.normal(size=shape)) gamma_rate = np.exp(np.random.normal(size=shape)) value = jnp.arange(15) num_samples = 300000 poisson_rate = np.random.gamma(gamma_conc, 1 / gamma_rate, size=(num_samples,) + shape) log_probs = dist.Poisson(poisson_rate).log_prob(value) expected = logsumexp(log_probs, 0) - jnp.log(num_samples) actual = dist.GammaPoisson(gamma_conc, gamma_rate).log_prob(value) assert_allclose(actual, expected, rtol=0.05)
Example #22
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 #23
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 5 votes |
def test_change_point_x64(): # Ref: https://forum.pyro.ai/t/i-dont-understand-why-nuts-code-is-not-working-bayesian-hackers-mail/696 warmup_steps, num_samples = 500, 3000 def model(data): alpha = 1 / jnp.mean(data) lambda1 = numpyro.sample('lambda1', dist.Exponential(alpha)) lambda2 = numpyro.sample('lambda2', dist.Exponential(alpha)) tau = numpyro.sample('tau', dist.Uniform(0, 1)) lambda12 = jnp.where(jnp.arange(len(data)) < tau * len(data), lambda1, lambda2) numpyro.sample('obs', dist.Poisson(lambda12), obs=data) count_data = jnp.array([ 13, 24, 8, 24, 7, 35, 14, 11, 15, 11, 22, 22, 11, 57, 11, 19, 29, 6, 19, 12, 22, 12, 18, 72, 32, 9, 7, 13, 19, 23, 27, 20, 6, 17, 13, 10, 14, 6, 16, 15, 7, 2, 15, 15, 19, 70, 49, 7, 53, 22, 21, 31, 19, 11, 18, 20, 12, 35, 17, 23, 17, 4, 2, 31, 30, 13, 27, 0, 39, 37, 5, 14, 13, 22, ]) kernel = NUTS(model=model) mcmc = MCMC(kernel, warmup_steps, num_samples) mcmc.run(random.PRNGKey(4), count_data) samples = mcmc.get_samples() tau_posterior = (samples['tau'] * len(count_data)).astype(jnp.int32) tau_values, counts = np.unique(tau_posterior, return_counts=True) mode_ind = jnp.argmax(counts) mode = tau_values[mode_ind] assert mode == 44 if 'JAX_ENABLE_X64' in os.environ: assert samples['lambda1'].dtype == jnp.float64 assert samples['lambda2'].dtype == jnp.float64 assert samples['tau'].dtype == jnp.float64
Example #24
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def test_beta_binomial_log_prob(total_count, shape): concentration0 = np.exp(np.random.normal(size=shape)) concentration1 = np.exp(np.random.normal(size=shape)) value = jnp.arange(1 + total_count) num_samples = 100000 probs = np.random.beta(concentration1, concentration0, size=(num_samples,) + shape) log_probs = dist.Binomial(total_count, probs).log_prob(value) expected = logsumexp(log_probs, 0) - jnp.log(num_samples) actual = dist.BetaBinomial(concentration1, concentration0, total_count).log_prob(value) assert_allclose(actual, expected, rtol=0.02)
Example #25
Source File: test_indexing.py From numpyro with Apache License 2.0 | 5 votes |
def test_hmm_example(prev_enum_dim, curr_enum_dim): hidden_dim = 8 probs_x = jnp.array(np.random.rand(hidden_dim, hidden_dim, hidden_dim)) x_prev = jnp.arange(hidden_dim).reshape((-1,) + (1,) * (-1 - prev_enum_dim)) x_curr = jnp.arange(hidden_dim).reshape((-1,) + (1,) * (-1 - curr_enum_dim)) expected = probs_x[x_prev.reshape(x_prev.shape + (1,)), x_curr.reshape(x_curr.shape + (1,)), jnp.arange(hidden_dim)] actual = Vindex(probs_x)[x_prev, x_curr, :] assert jnp.all(actual == expected)
Example #26
Source File: util.py From numpyro with Apache License 2.0 | 5 votes |
def cholesky_update(L, x, coef=1): """ Finds cholesky of L @ L.T + coef * x @ x.T. **References;** 1. A more efficient rank-one covariance matrix update for evolution strategies, Oswin Krause and Christian Igel """ batch_shape = lax.broadcast_shapes(L.shape[:-2], x.shape[:-1]) L = jnp.broadcast_to(L, batch_shape + L.shape[-2:]) x = jnp.broadcast_to(x, batch_shape + x.shape[-1:]) diag = jnp.diagonal(L, axis1=-2, axis2=-1) # convert to unit diagonal triangular matrix: L @ D @ T.t L = L / diag[..., None, :] D = jnp.square(diag) def scan_fn(carry, val): b, w = carry j, Dj, L_j = val wj = w[..., j] gamma = b * Dj + coef * jnp.square(wj) Dj_new = gamma / b b = gamma / Dj_new # update vectors w and L_j w = w - wj[..., None] * L_j L_j = L_j + (coef * wj / gamma)[..., None] * w return (b, w), (Dj_new, L_j) D, L = jnp.moveaxis(D, -1, 0), jnp.moveaxis(L, -1, 0) # move scan dim to front _, (D, L) = lax.scan(scan_fn, (jnp.ones(batch_shape), x), (jnp.arange(D.shape[0]), D, L)) D, L = jnp.moveaxis(D, 0, -1), jnp.moveaxis(L, 0, -1) # move scan dim back return L * jnp.sqrt(D)[..., None, :]
Example #27
Source File: ops.py From funsor with Apache License 2.0 | 5 votes |
def _new_arange(x, start, stop, step): return np.arange(start, stop, step)
Example #28
Source File: ops.py From funsor with Apache License 2.0 | 5 votes |
def _new_arange(x, stop): return np.arange(stop)
Example #29
Source File: space_serializer.py From trax with Apache License 2.0 | 5 votes |
def significance_map(self): return np.reshape(np.broadcast_to( np.arange(self._precision), self._space.shape + (self._precision,)), -1)
Example #30
Source File: resnet50.py From jaxnet with Apache License 2.0 | 5 votes |
def main(): key = PRNGKey(0) batch_size = 8 num_classes = 1001 input_shape = (224, 224, 3, batch_size) step_size = 0.1 num_steps = 10 resnet = ResNet50(num_classes) @parametrized def loss(inputs, targets): logits = resnet(inputs) return np.sum(logits * targets) @parametrized def accuracy(inputs, targets): target_class = np.argmax(targets, axis=-1) predicted_class = np.argmax(resnet(inputs), axis=-1) return np.mean(predicted_class == target_class) def synth_batches(): rng = npr.RandomState(0) while True: images = rng.rand(*input_shape).astype('float32') labels = rng.randint(num_classes, size=(batch_size, 1)) onehot_labels = labels == np.arange(num_classes) yield images, onehot_labels opt = optimizers.Momentum(step_size, mass=0.9) batches = synth_batches() print("\nInitializing parameters.") state = opt.init(loss.init_parameters(*next(batches), key=key)) for i in range(num_steps): print(f'Training on batch {i}.') state = opt.update(loss.apply, state, *next(batches)) trained_params = opt.get_parameters(state)