Python jax.numpy.array() Examples
The following are 30
code examples of jax.numpy.array().
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_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 #2
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 #3
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 #4
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def test_prior_with_sample_shape(): data = { "J": 8, "y": jnp.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]), "sigma": jnp.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]), } def schools_model(): mu = numpyro.sample('mu', dist.Normal(0, 5)) tau = numpyro.sample('tau', dist.HalfCauchy(5)) theta = numpyro.sample('theta', dist.Normal(mu, tau), sample_shape=(data['J'],)) numpyro.sample('obs', dist.Normal(theta, data['sigma']), obs=data['y']) num_samples = 500 mcmc = MCMC(NUTS(schools_model), num_warmup=500, num_samples=num_samples) mcmc.run(random.PRNGKey(0)) assert mcmc.get_samples()['theta'].shape == (num_samples, data['J'])
Example #5
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 #6
Source File: test_infer_util.py From numpyro with Apache License 2.0 | 6 votes |
def test_initialize_model_dirichlet_categorical(init_strategy): 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,)) rng_keys = random.split(random.PRNGKey(1), 2) init_params, _, _, _ = initialize_model(rng_keys, model, init_strategy=init_strategy, model_args=(data,)) for i in range(2): init_params_i, _, _, _ = initialize_model(rng_keys[i], model, init_strategy=init_strategy, model_args=(data,)) for name, p in init_params[0].items(): # XXX: the result is equal if we disable fast-math-mode assert_allclose(p[i], init_params_i[0][name], atol=1e-6)
Example #7
Source File: ucbadmit.py From numpyro with Apache License 2.0 | 6 votes |
def glmm(dept, male, applications, admit=None): v_mu = numpyro.sample('v_mu', dist.Normal(0, jnp.array([4., 1.]))) sigma = numpyro.sample('sigma', dist.HalfNormal(jnp.ones(2))) L_Rho = numpyro.sample('L_Rho', dist.LKJCholesky(2, concentration=2)) scale_tril = sigma[..., jnp.newaxis] * L_Rho # non-centered parameterization num_dept = len(jnp.unique(dept)) z = numpyro.sample('z', dist.Normal(jnp.zeros((num_dept, 2)), 1)) v = jnp.dot(scale_tril, z.T).T logits = v_mu[0] + v[dept, 0] + (v_mu[1] + v[dept, 1]) * male if admit is None: # we use a Delta site to record probs for predictive distribution probs = expit(logits) numpyro.sample('probs', dist.Delta(probs), obs=probs) numpyro.sample('admit', dist.Binomial(applications, logits=logits), obs=admit)
Example #8
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 #9
Source File: stochastic_volatility.py From numpyro with Apache License 2.0 | 6 votes |
def print_results(posterior, dates): def _print_row(values, row_name=''): quantiles = jnp.array([0.2, 0.4, 0.5, 0.6, 0.8]) row_name_fmt = '{:>8}' header_format = row_name_fmt + '{:>12}' * 5 row_format = row_name_fmt + '{:>12.3f}' * 5 columns = ['(p{})'.format(q * 100) for q in quantiles] q_values = jnp.quantile(values, quantiles, axis=0) print(header_format.format('', *columns)) print(row_format.format(row_name, *q_values)) print('\n') print('=' * 20, 'sigma', '=' * 20) _print_row(posterior['sigma']) print('=' * 20, 'nu', '=' * 20) _print_row(posterior['nu']) print('=' * 20, 'volatility', '=' * 20) for i in range(0, len(dates), 180): _print_row(jnp.exp(posterior['s'][:, i]), dates[i])
Example #10
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 #11
Source File: test_svi.py From numpyro with Apache License 2.0 | 6 votes |
def test_jitted_update_fn(): data = jnp.array([1.0] * 8 + [0.0] * 2) def model(data): f = numpyro.sample("beta", dist.Beta(1., 1.)) numpyro.sample("obs", dist.Bernoulli(f), obs=data) def guide(data): alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive) beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive) numpyro.sample("beta", dist.Beta(alpha_q, beta_q)) adam = optim.Adam(0.05) svi = SVI(model, guide, adam, ELBO()) svi_state = svi.init(random.PRNGKey(1), data) expected = svi.get_params(svi.update(svi_state, data)[0]) actual = svi.get_params(jit(svi.update)(svi_state, data=data)[0]) check_close(actual, expected, atol=1e-5)
Example #12
Source File: serialization_utils_test.py From trax with Apache License 2.0 | 6 votes |
def test_wrapped_policy_continuous(self, vocab_size): precision = 3 n_controls = 2 n_actions = 4 gin.bind_parameter('BoxSpaceSerializer.precision', precision) obs = np.array([[[1.5, 2], [-0.3, 1.23], [0.84, 0.07], [0.01, 0.66]]]) act = np.array([[[0, 1], [2, 0], [1, 3]]]) wrapped_policy = serialization_utils.wrap_policy( TestModel(extra_dim=vocab_size), # pylint: disable=no-value-for-parameter observation_space=gym.spaces.Box(shape=(2,), low=-2, high=2), action_space=gym.spaces.MultiDiscrete([n_actions] * n_controls), vocab_size=vocab_size, ) example = (obs, act) wrapped_policy.init(shapes.signature(example)) (act_logits, values) = wrapped_policy(example) self.assertEqual(act_logits.shape, obs.shape[:2] + (n_controls, n_actions)) self.assertEqual(values.shape, obs.shape[:2])
Example #13
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 #14
Source File: test_indexing.py From numpyro with Apache License 2.0 | 6 votes |
def test_value(x_shape, i_shape, j_shape, event_shape): x = jnp.array(np.random.rand(*(x_shape + (5, 6) + event_shape))) i = dist.Categorical(jnp.ones((5,))).sample(random.PRNGKey(1), i_shape) j = dist.Categorical(jnp.ones((6,))).sample(random.PRNGKey(2), j_shape) if event_shape: actual = Vindex(x)[..., i, j, :] else: actual = Vindex(x)[..., i, j] shape = lax.broadcast_shapes(x_shape, i_shape, j_shape) x = jnp.broadcast_to(x, shape + (5, 6) + event_shape) i = jnp.broadcast_to(i, shape) j = jnp.broadcast_to(j, shape) expected = np.empty(shape + event_shape, dtype=x.dtype) for ind in (itertools.product(*map(range, shape)) if shape else [()]): expected[ind] = x[ind + (i[ind].item(), j[ind].item())] assert jnp.all(actual == jnp.array(expected, dtype=x.dtype))
Example #15
Source File: tensor.py From discopy with BSD 3-Clause "New" or "Revised" License | 6 votes |
def __init__(self, *dim): """ >>> Id(1) Tensor(dom=Dim(1), cod=Dim(1), array=[1]) >>> list(Id(2).array.flatten()) [1.0, 0.0, 0.0, 1.0] >>> Id(2).array.shape (2, 2) >>> list(Id(2, 2).array.flatten())[:8] [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0] >>> list(Id(2, 2).array.flatten())[8:] [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0] """ dim = dim[0] if isinstance(dim[0], Dim) else Dim(*dim) array = functools.reduce( lambda a, x: np.tensordot(a, np.identity(x), 0) if a.shape else np.identity(x), dim, np.array(1)) array = np.moveaxis( array, [2 * i for i in range(len(dim))], list(range(len(dim)))) super().__init__(dim, dim, array)
Example #16
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 #17
Source File: scan.py From numpyro with Apache License 2.0 | 6 votes |
def scan_wrapper(f, init, xs, length, reverse, rng_key=None, substitute_stack=[]): def body_fn(wrapped_carry, x): i, rng_key, carry = wrapped_carry rng_key, subkey = random.split(rng_key) if rng_key is not None else (None, None) with handlers.block(): seeded_fn = handlers.seed(f, subkey) if subkey is not None else f for subs_type, subs_map in substitute_stack: subs_fn = partial(_subs_wrapper, subs_map, i, length) if subs_type == 'condition': seeded_fn = handlers.condition(seeded_fn, condition_fn=subs_fn) elif subs_type == 'substitute': seeded_fn = handlers.substitute(seeded_fn, substitute_fn=subs_fn) with handlers.trace() as trace: carry, y = seeded_fn(carry, x) return (i + 1, rng_key, carry), (PytreeTrace(trace), y) if length is None: length = tree_flatten(xs)[0][0].shape[0] return lax.scan(body_fn, (jnp.array(0), rng_key, init), xs, length=length, reverse=reverse)
Example #18
Source File: test_distributions.py From numpyro with Apache License 2.0 | 6 votes |
def test_ZIP_log_prob(rate): # if gate is 0 ZIP is Poisson zip_ = dist.ZeroInflatedPoisson(0., rate) pois = dist.Poisson(rate) s = zip_.sample(random.PRNGKey(0), (20,)) zip_prob = zip_.log_prob(s) pois_prob = pois.log_prob(s) assert_allclose(zip_prob, pois_prob) # if gate is 1 ZIP is Delta(0) zip_ = dist.ZeroInflatedPoisson(1., rate) delta = dist.Delta(0.) s = jnp.array([0., 1.]) zip_prob = zip_.log_prob(s) delta_prob = delta.log_prob(s) assert_allclose(zip_prob, delta_prob)
Example #19
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def stirling_approx_tail(k): precomputed = jnp.array([ 0.08106146679532726, 0.04134069595540929, 0.02767792568499834, 0.02079067210376509, 0.01664469118982119, 0.01387612882307075, 0.01189670994589177, 0.01041126526197209, 0.009255462182712733, 0.008330563433362871, ]) kp1 = k + 1 kp1sq = (k + 1) ** 2 return jnp.where(k < 10, precomputed[k], (1. / 12 - (1. / 360 - (1. / 1260) / kp1sq) / kp1sq) / kp1)
Example #20
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 #21
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 #22
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def gen_values_outside_bounds(constraint, size, key=random.PRNGKey(11)): if isinstance(constraint, constraints._Boolean): return random.bernoulli(key, shape=size) - 2 elif isinstance(constraint, constraints._GreaterThan): return constraint.lower_bound - jnp.exp(random.normal(key, size)) elif isinstance(constraint, constraints._IntegerInterval): lower_bound = jnp.broadcast_to(constraint.lower_bound, size) return random.randint(key, size, lower_bound - 1, lower_bound) elif isinstance(constraint, constraints._IntegerGreaterThan): return constraint.lower_bound - random.poisson(key, np.array(5), shape=size) elif isinstance(constraint, constraints._Interval): upper_bound = jnp.broadcast_to(constraint.upper_bound, size) return random.uniform(key, size, minval=upper_bound, maxval=upper_bound + 1.) elif isinstance(constraint, (constraints._Real, constraints._RealVector)): return lax.full(size, jnp.nan) elif isinstance(constraint, constraints._Simplex): return osp.dirichlet.rvs(alpha=jnp.ones((size[-1],)), size=size[:-1]) + 1e-2 elif isinstance(constraint, constraints._Multinomial): n = size[-1] return multinomial(key, p=jnp.ones((n,)) / n, n=constraint.upper_bound, shape=size[:-1]) + 1 elif isinstance(constraint, constraints._CorrCholesky): return signed_stick_breaking_tril( random.uniform(key, size[:-2] + (size[-1] * (size[-1] - 1) // 2,), minval=-1, maxval=1)) + 1e-2 elif isinstance(constraint, constraints._CorrMatrix): cholesky = 1e-2 + signed_stick_breaking_tril( random.uniform(key, size[:-2] + (size[-1] * (size[-1] - 1) // 2,), minval=-1, maxval=1)) return jnp.matmul(cholesky, jnp.swapaxes(cholesky, -2, -1)) elif isinstance(constraint, constraints._LowerCholesky): return random.uniform(key, size) elif isinstance(constraint, constraints._PositiveDefinite): return random.normal(key, size) elif isinstance(constraint, constraints._OrderedVector): x = jnp.cumsum(random.exponential(key, size), -1) return x[..., ::-1] else: raise NotImplementedError('{} not implemented.'.format(constraint))
Example #23
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def gen_values_within_bounds(constraint, size, key=random.PRNGKey(11)): eps = 1e-6 if isinstance(constraint, constraints._Boolean): return random.bernoulli(key, shape=size) elif isinstance(constraint, constraints._GreaterThan): return jnp.exp(random.normal(key, size)) + constraint.lower_bound + eps elif isinstance(constraint, constraints._IntegerInterval): lower_bound = jnp.broadcast_to(constraint.lower_bound, size) upper_bound = jnp.broadcast_to(constraint.upper_bound, size) return random.randint(key, size, lower_bound, upper_bound + 1) elif isinstance(constraint, constraints._IntegerGreaterThan): return constraint.lower_bound + random.poisson(key, np.array(5), shape=size) elif isinstance(constraint, constraints._Interval): lower_bound = jnp.broadcast_to(constraint.lower_bound, size) upper_bound = jnp.broadcast_to(constraint.upper_bound, size) return random.uniform(key, size, minval=lower_bound, maxval=upper_bound) elif isinstance(constraint, (constraints._Real, constraints._RealVector)): return random.normal(key, size) elif isinstance(constraint, constraints._Simplex): return osp.dirichlet.rvs(alpha=jnp.ones((size[-1],)), size=size[:-1]) elif isinstance(constraint, constraints._Multinomial): n = size[-1] return multinomial(key, p=jnp.ones((n,)) / n, n=constraint.upper_bound, shape=size[:-1]) elif isinstance(constraint, constraints._CorrCholesky): return signed_stick_breaking_tril( random.uniform(key, size[:-2] + (size[-1] * (size[-1] - 1) // 2,), minval=-1, maxval=1)) elif isinstance(constraint, constraints._CorrMatrix): cholesky = signed_stick_breaking_tril( random.uniform(key, size[:-2] + (size[-1] * (size[-1] - 1) // 2,), minval=-1, maxval=1)) return jnp.matmul(cholesky, jnp.swapaxes(cholesky, -2, -1)) elif isinstance(constraint, constraints._LowerCholesky): return jnp.tril(random.uniform(key, size)) elif isinstance(constraint, constraints._PositiveDefinite): x = random.normal(key, size) return jnp.matmul(x, jnp.swapaxes(x, -2, -1)) elif isinstance(constraint, constraints._OrderedVector): x = jnp.cumsum(random.exponential(key, size), -1) return x - random.normal(key, size[:-1]) else: raise NotImplementedError('{} not implemented.'.format(constraint))
Example #24
Source File: test_infer_util.py From numpyro with Apache License 2.0 | 5 votes |
def test_model_with_transformed_distribution(): x_prior = dist.HalfNormal(2) y_prior = dist.LogNormal(scale=3.) # transformed distribution def model(): numpyro.sample('x', x_prior) numpyro.sample('y', y_prior) params = {'x': jnp.array(-5.), 'y': jnp.array(7.)} model = handlers.seed(model, random.PRNGKey(0)) inv_transforms = {'x': biject_to(x_prior.support), 'y': biject_to(y_prior.support)} expected_samples = partial(transform_fn, inv_transforms)(params) expected_potential_energy = ( - x_prior.log_prob(expected_samples['x']) - y_prior.log_prob(expected_samples['y']) - inv_transforms['x'].log_abs_det_jacobian(params['x'], expected_samples['x']) - inv_transforms['y'].log_abs_det_jacobian(params['y'], expected_samples['y']) ) reparam_model = handlers.reparam(model, {'y': TransformReparam()}) base_params = {'x': params['x'], 'y_base': params['y']} actual_samples = constrain_fn( handlers.seed(reparam_model, random.PRNGKey(0)), (), {}, base_params, return_deterministic=True) actual_potential_energy = potential_energy(reparam_model, (), {}, base_params) assert_allclose(expected_samples['x'], actual_samples['x']) assert_allclose(expected_samples['y'], actual_samples['y']) assert_allclose(actual_potential_energy, expected_potential_energy)
Example #25
Source File: test_mcmc.py From numpyro with Apache License 2.0 | 5 votes |
def test_chain_smoke(chain_method, compile_args): 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 data = dist.Categorical(jnp.array([0.1, 0.6, 0.3])).sample(random.PRNGKey(1), (2000,)) kernel = NUTS(model) mcmc = MCMC(kernel, 2, 5, num_chains=2, chain_method=chain_method, jit_model_args=compile_args) mcmc.warmup(random.PRNGKey(0), data) mcmc.run(random.PRNGKey(1), data)
Example #26
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 #27
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 #28
Source File: auto_reg_nn.py From numpyro with Apache License 2.0 | 5 votes |
def create_mask(input_dim, hidden_dims, permutation, output_dim_multiplier): """ Creates (non-conditional) MADE masks :param input_dim: the dimensionality of the input variable :type input_dim: int :param hidden_dims: the dimensionality of the hidden layers(s) :type hidden_dims: list[int] :param permutation: the order of the input variables :type permutation: numpy array of integers of length `input_dim` :param output_dim_multiplier: tiles the output (e.g. for when a separate mean and scale parameter are desired) :type output_dim_multiplier: int """ # Create mask indices for input, hidden layers, and final layer var_index = jnp.zeros(permutation.shape[0]) var_index = ops.index_update(var_index, permutation, jnp.arange(input_dim)) # Create the indices that are assigned to the neurons input_indices = 1 + var_index hidden_indices = [sample_mask_indices(input_dim - 1, h) for h in hidden_dims] output_indices = jnp.tile(var_index + 1, output_dim_multiplier) # Create mask from input to output for the skips connections mask_skip = output_indices[None, :] > input_indices[:, None] # Create mask from input to first hidden layer, and between subsequent hidden layers # NB: these masks are transposed version of Pyro ones masks = [hidden_indices[0][None, :] >= input_indices[:, None]] for i in range(1, len(hidden_dims)): masks.append(hidden_indices[i][None, :] >= hidden_indices[i - 1][:, None]) # Create mask from last hidden layer to output layer masks.append(output_indices[None, :] > hidden_indices[-1][:, None]) return masks, mask_skip
Example #29
Source File: optim.py From numpyro with Apache License 2.0 | 5 votes |
def init(self, params: _Params) -> _IterOptState: """ Initialize the optimizer with parameters designated to be optimized. :param params: a collection of numpy arrays. :return: initial optimizer state. """ opt_state = self.init_fn(params) return jnp.array(0), opt_state
Example #30
Source File: autoguide.py From numpyro with Apache License 2.0 | 5 votes |
def quantiles(self, params, quantiles): transform = self.get_transform(params) quantiles = jnp.array(quantiles)[..., None] latent = dist.Normal(transform.loc, jnp.diagonal(transform.scale_tril)).icdf(quantiles) return self._unpack_and_constrain(latent, params)