Python jax.numpy.sqrt() Examples
The following are 30
code examples of jax.numpy.sqrt().
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: modules.py From jaxnet with Apache License 2.0 | 6 votes |
def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True, beta_init=zeros, gamma_init=ones): """Layer construction function for a batch normalization layer.""" axis = (axis,) if np.isscalar(axis) else axis @parametrized def batch_norm(x): ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x))) mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True) z = (x - mean) / np.sqrt(var + epsilon) shape = tuple(d for i, d in enumerate(x.shape) if i not in axis) scaled = z * parameter(shape, gamma_init, 'gamma')[ed] if scale else z return scaled + parameter(shape, beta_init, 'beta')[ed] if center else scaled return batch_norm
Example #2
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 #3
Source File: continuous.py From numpyro with Apache License 2.0 | 6 votes |
def _onion(self, key, size): key_beta, key_normal = random.split(key) # Now we generate w term in Algorithm 3.2 of [1]. beta_sample = self._beta.sample(key_beta, size) # The following Normal distribution is used to create a uniform distribution on # a hypershere (ref: http://mathworld.wolfram.com/HyperspherePointPicking.html) normal_sample = random.normal( key_normal, shape=size + self.batch_shape + (self.dimension * (self.dimension - 1) // 2,) ) normal_sample = vec_to_tril_matrix(normal_sample, diagonal=0) u_hypershere = normal_sample / jnp.linalg.norm(normal_sample, axis=-1, keepdims=True) w = jnp.expand_dims(jnp.sqrt(beta_sample), axis=-1) * u_hypershere # put w into the off-diagonal triangular part cholesky = ops.index_add(jnp.zeros(size + self.batch_shape + self.event_shape), ops.index[..., 1:, :-1], w) # correct the diagonal # NB: we clip due to numerical precision diag = jnp.sqrt(jnp.clip(1 - jnp.sum(cholesky ** 2, axis=-1), a_min=0.)) cholesky = cholesky + jnp.expand_dims(diag, axis=-1) * jnp.identity(self.dimension) return cholesky
Example #4
Source File: continuous.py From numpyro with Apache License 2.0 | 6 votes |
def _cvine(self, key, size): # C-vine method first uses beta_dist to generate partial correlations, # then apply signed stick breaking to transform to cholesky factor. # Here is an attempt to prove that using signed stick breaking to # generate correlation matrices is the same as the C-vine method in [1] # for the entry r_32. # # With notations follow from [1], we define # p: partial correlation matrix, # c: cholesky factor, # r: correlation matrix. # From recursive formula (2) in [1], we have # r_32 = p_32 * sqrt{(1 - p_21^2)*(1 - p_31^2)} + p_21 * p_31 =: I # On the other hand, signed stick breaking process gives: # l_21 = p_21, l_31 = p_31, l_22 = sqrt(1 - p_21^2), l_32 = p_32 * sqrt(1 - p_31^2) # r_32 = l_21 * l_31 + l_22 * l_32 # = p_21 * p_31 + p_32 * sqrt{(1 - p_21^2)*(1 - p_31^2)} = I beta_sample = self._beta.sample(key, size) partial_correlation = 2 * beta_sample - 1 # scale to domain to (-1, 1) return signed_stick_breaking_tril(partial_correlation)
Example #5
Source File: util.py From numpyro with Apache License 2.0 | 6 votes |
def _get_tr_params(n, p): # See Table 1. Additionally, we pre-compute log(p), log1(-p) and the # constant terms, that depend only on (n, p, m) in log(f(k)) (bottom of page 5). mu = n * p spq = jnp.sqrt(mu * (1 - p)) c = mu + 0.5 b = 1.15 + 2.53 * spq a = -0.0873 + 0.0248 * b + 0.01 * p alpha = (2.83 + 5.1 / b) * spq u_r = 0.43 v_r = 0.92 - 4.2 / b m = jnp.floor((n + 1) * p).astype(n.dtype) log_p = jnp.log(p) log1_p = jnp.log1p(-p) log_h = (m + 0.5) * (jnp.log((m + 1.) / (n - m + 1.)) + log1_p - log_p) + \ (stirling_approx_tail(m) + stirling_approx_tail(n - m)) return _tr_params(c, b, a, alpha, u_r, v_r, m, log_p, log1_p, log_h)
Example #6
Source File: mcmc.py From numpyro with Apache License 2.0 | 6 votes |
def _get_proposal_loc_and_scale(samples, loc, scale, new_sample): # get loc/scale of q_{-n} (Algorithm 1, line 5 of ref [1]) for n from 1 -> N # these loc/scale will be stacked to the first dim; so # proposal_loc.shape[0] = proposal_loc.shape[0] = N # Here, we use the numerical stability procedure in Appendix 6 of [1]. weight = 1 / samples.shape[0] if scale.ndim > loc.ndim: new_scale = cholesky_update(scale, new_sample - loc, weight) proposal_scale = cholesky_update(new_scale, samples - loc, -weight) proposal_scale = cholesky_update(proposal_scale, new_sample - samples, - (weight ** 2)) else: var = jnp.square(scale) + weight * jnp.square(new_sample - loc) proposal_var = var - weight * jnp.square(samples - loc) proposal_var = proposal_var - weight ** 2 * jnp.square(new_sample - samples) proposal_scale = jnp.sqrt(proposal_var) proposal_loc = loc + weight * (new_sample - samples) return proposal_loc, proposal_scale
Example #7
Source File: bnn.py From numpyro with Apache License 2.0 | 6 votes |
def model(X, Y, D_H): D_X, D_Y = X.shape[1], 1 # sample first layer (we put unit normal priors on all weights) w1 = numpyro.sample("w1", dist.Normal(jnp.zeros((D_X, D_H)), jnp.ones((D_X, D_H)))) # D_X D_H z1 = nonlin(jnp.matmul(X, w1)) # N D_H <= first layer of activations # sample second layer w2 = numpyro.sample("w2", dist.Normal(jnp.zeros((D_H, D_H)), jnp.ones((D_H, D_H)))) # D_H D_H z2 = nonlin(jnp.matmul(z1, w2)) # N D_H <= second layer of activations # sample final layer of weights and neural network output w3 = numpyro.sample("w3", dist.Normal(jnp.zeros((D_H, D_Y)), jnp.ones((D_H, D_Y)))) # D_H D_Y z3 = jnp.matmul(z2, w3) # N D_Y <= output of the neural network # we put a prior on the observation noise prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0)) sigma_obs = 1.0 / jnp.sqrt(prec_obs) # observe data numpyro.sample("Y", dist.Normal(z3, sigma_obs), obs=Y) # helper function for HMC inference
Example #8
Source File: pixelcnn.py From jaxnet with Apache License 2.0 | 6 votes |
def ConvOrConvTranspose(out_chan, filter_shape=(3, 3), strides=None, padding='SAME', init_scale=1., transpose=False): strides = strides or (1,) * len(filter_shape) def apply(inputs, V, g, b): V = g * _l2_normalize(V, (0, 1, 2)) return (lax.conv_transpose if transpose else _conv)(inputs, V, strides, padding) - b @parametrized def conv_or_conv_transpose(inputs): V = parameter(filter_shape + (inputs.shape[-1], out_chan), normal(.05), 'V') example_out = apply(inputs, V=V, g=jnp.ones(out_chan), b=jnp.zeros(out_chan)) # TODO remove need for `.aval.val` when capturing variables in initializer function: g = Parameter(lambda key: init_scale / jnp.sqrt(jnp.var(example_out.aval.val, (0, 1, 2)) + 1e-10), 'g')() b = Parameter(lambda key: jnp.mean(example_out.aval.val, (0, 1, 2)) * g.aval.val, 'b')() return apply(inputs, V, b, g) return conv_or_conv_transpose
Example #9
Source File: jax_backend.py From pyhf with Apache License 2.0 | 5 votes |
def normal_logpdf(self, x, mu, sigma): # this is much faster than # norm.logpdf(x, loc=mu, scale=sigma) # https://codereview.stackexchange.com/questions/69718/fastest-computation-of-n-likelihoods-on-normal-distributions root2 = np.sqrt(2) root2pi = np.sqrt(2 * np.pi) prefactor = -np.log(sigma * root2pi) summand = -np.square(np.divide((x - mu), (root2 * sigma))) return prefactor + summand # def normal_logpdf(self, x, mu, sigma): # return norm.logpdf(x, loc=mu, scale=sigma)
Example #10
Source File: ops.py From funsor with Apache License 2.0 | 5 votes |
def _cholesky(x): """ Like :func:`numpy.linalg.cholesky` but uses sqrt for scalar matrices. """ if x.shape[-1] == 1: return np.sqrt(x) return np.linalg.cholesky(x)
Example #11
Source File: test_distributions.py From numpyro with Apache License 2.0 | 5 votes |
def _tril_cholesky_to_tril_corr(x): w = vec_to_tril_matrix(x, diagonal=-1) diag = jnp.sqrt(1 - jnp.sum(w ** 2, axis=-1)) cholesky = w + jnp.expand_dims(diag, axis=-1) * jnp.identity(w.shape[-1]) corr = jnp.matmul(cholesky, cholesky.T) return matrix_to_tril_vec(corr, diagonal=-1)
Example #12
Source File: test_hmc_util.py From numpyro with Apache License 2.0 | 5 votes |
def test_welford_covariance(jitted, diagonal, regularize): with optional(jitted, disable_jit()), optional(jitted, control_flow_prims_disabled()): np.random.seed(0) loc = np.random.randn(3) a = np.random.randn(3, 3) target_cov = np.matmul(a, a.T) x = np.random.multivariate_normal(loc, target_cov, size=(2000,)) x = device_put(x) @jit def get_cov(x): wc_init, wc_update, wc_final = welford_covariance(diagonal=diagonal) wc_state = wc_init(3) wc_state = fori_loop(0, 2000, lambda i, val: wc_update(x[i], val), wc_state) cov, cov_inv_sqrt = wc_final(wc_state, regularize=regularize) return cov, cov_inv_sqrt cov, cov_inv_sqrt = get_cov(x) if diagonal: diag_cov = jnp.diagonal(target_cov) assert_allclose(cov, diag_cov, rtol=0.06) assert_allclose(cov_inv_sqrt, jnp.sqrt(jnp.reciprocal(diag_cov)), rtol=0.06) else: assert_allclose(cov, target_cov, rtol=0.06) assert_allclose(cov_inv_sqrt, jnp.linalg.cholesky(jnp.linalg.inv(cov)), rtol=0.06) ######################################## # verlocity_verlet Test ########################################
Example #13
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def sample(self, key, sample_shape=()): key_normal, key_chi2 = random.split(key) std_normal = random.normal(key_normal, shape=sample_shape + self.batch_shape) z = self._chi2.sample(key_chi2, sample_shape) y = std_normal * jnp.sqrt(self.df / z) return self.loc + self.scale * y
Example #14
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def log_prob(self, value): normalize_term = jnp.log(jnp.sqrt(2 * jnp.pi) * self.scale) value_scaled = (value - self.loc) / self.scale return -0.5 * value_scaled ** 2 - normalize_term
Example #15
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def scale_tril(self): # The following identity is used to increase the numerically computation stability # for Cholesky decomposition (see http://www.gaussianprocess.org/gpml/, Section 3.4.3): # W @ W.T + D = D1/2 @ (I + D-1/2 @ W @ W.T @ D-1/2) @ D1/2 # The matrix "I + D-1/2 @ W @ W.T @ D-1/2" has eigenvalues bounded from below by 1, # hence it is well-conditioned and safe to take Cholesky decomposition. cov_diag_sqrt_unsqueeze = jnp.expand_dims(jnp.sqrt(self.cov_diag), axis=-1) Dinvsqrt_W = self.cov_factor / cov_diag_sqrt_unsqueeze K = jnp.matmul(Dinvsqrt_W, jnp.swapaxes(Dinvsqrt_W, -1, -2)) K = jnp.add(K, jnp.identity(K.shape[-1])) scale_tril = cov_diag_sqrt_unsqueeze * jnp.linalg.cholesky(K) return scale_tril
Example #16
Source File: mnist_vae.py From jaxnet with Apache License 2.0 | 5 votes |
def gaussian_sample(key, mu, sigmasq): """Sample a diagonal Gaussian.""" return mu + np.sqrt(sigmasq) * random.normal(key, mu.shape)
Example #17
Source File: pixelcnn.py From jaxnet with Apache License 2.0 | 5 votes |
def _l2_normalize(arr, axis): return arr / jnp.sqrt(jnp.sum(arr ** 2, axis=axis, keepdims=True))
Example #18
Source File: continuous.py From numpyro with Apache License 2.0 | 5 votes |
def mean(self): return jnp.sqrt(2 / jnp.pi) * self.scale
Example #19
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 #20
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 #21
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 #22
Source File: transforms.py From numpyro with Apache License 2.0 | 5 votes |
def __call__(self, x): n = round((math.sqrt(1 + 8 * x.shape[-1]) - 1) / 2) z = vec_to_tril_matrix(x[..., :-n], diagonal=-1) diag = jnp.exp(x[..., -n:]) return z + jnp.expand_dims(diag, axis=-1) * jnp.identity(n)
Example #23
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 #24
Source File: jax.py From deepx with MIT License | 5 votes |
def sqrt(self, x): return np.sqrt(x)
Example #25
Source File: gp.py From numpyro with Apache License 2.0 | 5 votes |
def predict(rng_key, X, Y, X_test, var, length, noise): # compute kernels between train and test data, etc. k_pp = kernel(X_test, X_test, var, length, noise, include_noise=True) k_pX = kernel(X_test, X, var, length, noise, include_noise=False) k_XX = kernel(X, X, var, length, noise, include_noise=True) K_xx_inv = jnp.linalg.inv(k_XX) K = k_pp - jnp.matmul(k_pX, jnp.matmul(K_xx_inv, jnp.transpose(k_pX))) sigma_noise = jnp.sqrt(jnp.clip(jnp.diag(K), a_min=0.)) * jax.random.normal(rng_key, X_test.shape[:1]) mean = jnp.matmul(k_pX, jnp.matmul(K_xx_inv, Y)) # we return both the mean function and a sample from the posterior predictive for the # given set of hyperparameters return mean, mean + sigma_noise # create artificial regression dataset
Example #26
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 #27
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 #28
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 5 votes |
def compute_pairwise_mean_variance(X, Y, dim1, dim2, msq, lam, eta1, xisq, c, var_obs): P, N = X.shape[1], X.shape[0] probe = jnp.zeros((4, P)) probe = jax.ops.index_update(probe, jax.ops.index[:, dim1], jnp.array([1.0, 1.0, -1.0, -1.0])) probe = jax.ops.index_update(probe, jax.ops.index[:, dim2], jnp.array([1.0, -1.0, 1.0, -1.0])) eta2 = jnp.square(eta1) * jnp.sqrt(xisq) / msq kappa = jnp.sqrt(msq) * lam / jnp.sqrt(msq + jnp.square(eta1 * lam)) kX = kappa * X kprobe = kappa * probe k_xx = kernel(kX, kX, eta1, eta2, c) + var_obs * jnp.eye(N) k_xx_inv = jnp.linalg.inv(k_xx) k_probeX = kernel(kprobe, kX, eta1, eta2, c) k_prbprb = kernel(kprobe, kprobe, eta1, eta2, c) vec = jnp.array([0.25, -0.25, -0.25, 0.25]) mu = jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, Y)) mu = jnp.dot(mu, vec) var = k_prbprb - jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, jnp.transpose(k_probeX))) var = jnp.matmul(var, vec) var = jnp.dot(var, vec) return mu, var # Sample coefficients theta from the posterior for a given MCMC sample. # The first P returned values are {theta_1, theta_2, ...., theta_P}, while # the remaining values are {theta_ij} for i,j in the list `active_dims`, # sorted so that i < j.
Example #29
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 5 votes |
def compute_singleton_mean_variance(X, Y, dimension, msq, lam, eta1, xisq, c, var_obs): P, N = X.shape[1], X.shape[0] probe = jnp.zeros((2, P)) probe = jax.ops.index_update(probe, jax.ops.index[:, dimension], jnp.array([1.0, -1.0])) eta2 = jnp.square(eta1) * jnp.sqrt(xisq) / msq kappa = jnp.sqrt(msq) * lam / jnp.sqrt(msq + jnp.square(eta1 * lam)) kX = kappa * X kprobe = kappa * probe k_xx = kernel(kX, kX, eta1, eta2, c) + var_obs * jnp.eye(N) k_xx_inv = jnp.linalg.inv(k_xx) k_probeX = kernel(kprobe, kX, eta1, eta2, c) k_prbprb = kernel(kprobe, kprobe, eta1, eta2, c) vec = jnp.array([0.50, -0.50]) mu = jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, Y)) mu = jnp.dot(mu, vec) var = k_prbprb - jnp.matmul(k_probeX, jnp.matmul(k_xx_inv, jnp.transpose(k_probeX))) var = jnp.matmul(var, vec) var = jnp.dot(var, vec) return mu, var # Compute the mean and variance of coefficient theta_ij for a MCMC sample of the # kernel hyperparameters (eta1, xisq, ...). Compare to theorem 5.1 in reference [1].
Example #30
Source File: sparse_regression.py From numpyro with Apache License 2.0 | 5 votes |
def model(X, Y, hypers): S, P, N = hypers['expected_sparsity'], X.shape[1], X.shape[0] sigma = numpyro.sample("sigma", dist.HalfNormal(hypers['alpha3'])) phi = sigma * (S / jnp.sqrt(N)) / (P - S) eta1 = numpyro.sample("eta1", dist.HalfCauchy(phi)) msq = numpyro.sample("msq", dist.InverseGamma(hypers['alpha1'], hypers['beta1'])) xisq = numpyro.sample("xisq", dist.InverseGamma(hypers['alpha2'], hypers['beta2'])) eta2 = jnp.square(eta1) * jnp.sqrt(xisq) / msq lam = numpyro.sample("lambda", dist.HalfCauchy(jnp.ones(P))) kappa = jnp.sqrt(msq) * lam / jnp.sqrt(msq + jnp.square(eta1 * lam)) # sample observation noise var_obs = numpyro.sample("var_obs", dist.InverseGamma(hypers['alpha_obs'], hypers['beta_obs'])) # compute kernel kX = kappa * X k = kernel(kX, kX, eta1, eta2, hypers['c']) + var_obs * jnp.eye(N) assert k.shape == (N, N) # sample Y according to the standard gaussian process formula numpyro.sample("Y", dist.MultivariateNormal(loc=jnp.zeros(X.shape[0]), covariance_matrix=k), obs=Y) # Compute the mean and variance of coefficient theta_i (where i = dimension) for a # MCMC sample of the kernel hyperparameters (eta1, xisq, ...). # Compare to theorem 5.1 in reference [1].