AbstractIntegrator |
class |
Protocol for Gaussian integral approximation. |
AbstractLikelihood |
class |
Base class for likelihood functions with optional analytical ELL. |
AbstractLogdetStrategy |
class |
Protocol for log-determinant strategies. |
AbstractPreconditioner |
class |
Protocol for preconditioners producing an approximate inverse. |
AbstractSolveStrategy |
class |
Protocol for linear solve strategies. |
AbstractSolverStrategy |
class |
Protocol for solver strategies that pair solve + logdet. |
AnalyticalPsiStatistics |
class |
Protocol for kernels with closed-form Ψ statistics. |
AssumedDensityFilter |
class |
KL-optimal Gaussian projection via moment matching. |
AutoSolver |
class |
Automatic solver selection based on operator type and size. |
BBMMSolver |
class |
Black-Box Matrix-Matrix solver (Gardner et al. 2018). |
BernoulliLikelihood |
class |
Bernoulli likelihood with logit link. |
BlockDiag |
class |
Block diagonal operator diag(A₁, A₂, …, Aₖ). |
BlockTriDiag |
class |
Symmetric block-tridiagonal operator. |
CGSolver |
class |
Iterative CG solver with stochastic log-determinant. |
CapacitanceSolver |
class |
Solve a base system subject to homogeneous point constraints. |
Circulant |
function |
(Block-)circulant operator from its first column. |
ComposedSolver |
class |
Mix-and-match solve and logdet from different strategies. |
ConstantSDE |
class |
State-space representation of a constant kernel. |
CosineSDE |
class |
State-space representation of the cosine kernel. |
CubatureIntegrator |
class |
Third-order spherical-radial cubature. |
DAREResult |
class |
Result of DARE solver. |
DenseFallbackWarning |
class |
Warning emitted when a structured primitive materialises an operator. |
DenseLogdet |
class |
Dense log-determinant via gaussx structural dispatch. |
DenseSolver |
class |
Dense solver strategy using gaussx structural dispatch. |
DiagonalisedOperator |
class |
Linear operator A x = V⁻¹(Λ ⊙ V x) given by a transform pair. |
EigenFactorization |
class |
Precomputed eigendecomposition A = V diag(λ) V⁻¹ of a square matrix. |
EigenProPreconditioner |
class |
Spectral correction state for EigenPro kernel SGD. |
EmissionModel |
class |
Observation (emission) model wrapping a linear observation matrix. |
FalkonPreconditioner |
class |
The two upper-triangular Cholesky factors of the Falkon preconditioner. |
FifthOrderCubatureIntegrator |
class |
Fifth-order fully symmetric cubature (McNamee & Stenger). |
FilterState |
class |
Output of kalman_filter. |
GaussHermiteIntegrator |
class |
Gauss-Hermite quadrature integrator. |
GaussianExpFam |
class |
Gaussian in natural (exponential family) parameters. |
GaussianLikelihood |
class |
Gaussian likelihood log N(y \| f, noise_var * I). |
GaussianSites |
class |
Time-varying Gaussian likelihood sites in natural parameterization. |
GaussianState |
class |
Gaussian distribution as (mean, covariance operator) pair. |
HeteroscedasticGaussianLikelihood |
class |
Heteroscedastic Gaussian likelihood with input-dependent noise. |
ImplicitCrossKernelOperator |
class |
Matrix-free rectangular kernel operator K(X, Z) \cdot v. |
ImplicitKernelOperator |
class |
Matrix-free kernel operator: (K + sigma^2 I) v via sequential scan. |
IndefiniteSLQLogdet |
class |
Stochastic log\|det(A)\| for symmetric (possibly indefinite) operators. |
InfiniteHorizonState |
class |
Output of infinite_horizon_filter. |
IntegratedWienerSDE |
class |
State-space representation of an integrated Wiener process. |
InterpolatedOperator |
class |
Structured Kernel Interpolation: K \approx W K_{uu} W^T. |
JacobiPreconditioner |
class |
Diagonal preconditioner M^{-1} = diag(1 / diag(A)). |
KernelOperator |
class |
Kernel matrix operator with efficient first-order autodiff. |
Kronecker |
class |
Kronecker product operator A₁ ⊗ A₂ ⊗ … ⊗ Aₖ. |
KroneckerSum |
class |
Kronecker sum A \oplus B = A \otimes I_b + I_a \otimes B. |
KroneckerSumSqrt |
class |
Symmetric square root of A \oplus B via per-factor eigenvectors. |
LOOResult |
class |
Result of leave-one-out cross-validation. |
LOVECache |
class |
Cached Lanczos factorization for fast predictive variance. |
LSMRSolver |
class |
LSMR iterative least-squares solver (Fong & Saunders 2011). |
LowRankUpdate |
class |
Low-rank update operator L + U diag(d) Vᵀ. |
LowerBlockTriDiag |
class |
Lower triangular block-bidiagonal Cholesky factor. |
MINRESSolver |
class |
MINRES solver for symmetric (possibly indefinite) systems. |
MaskedOperator |
class |
Row/column-masked view of a base operator. |
MaternSDE |
class |
State-space representation of the Matern kernel. |
MomentMatchResult |
class |
Tilted log-normaliser and its derivatives w.r.t. the cavity mean. |
MonteCarloIntegrator |
class |
Monte Carlo moment matching: sample, propagate, compute moments. |
NystromPreconditioner |
class |
Low-rank approximate inverse from randomized operator probing. |
OperatorPreconditioner |
class |
Use an externally supplied approximate inverse as a preconditioner. |
PartialCholeskyPreconditioner |
class |
Preconditioner from a pivoted partial Cholesky factor. |
PeriodicSDE |
class |
State-space representation of the periodic (MacKay) kernel. |
PoissonLikelihood |
class |
Poisson likelihood with log link. |
PreconditionedCGSolver |
class |
CG solver with pivoted partial Cholesky preconditioner. |
PredictionCache |
class |
Cached training solve for amortized predictions. |
ProductOperator |
function |
Lazy matmul (A B) v = A (B v). |
ProductSDE |
class |
Product of two SDE kernels via Kronecker composition. |
PropagationResult |
class |
Output of uncertainty propagation through a nonlinear function. |
QuasiPeriodicSDE |
class |
Quasi-periodic kernel: product of Matern and Periodic SDE. |
RootDecomposition |
class |
Tall factor R with R Rᵀ ≈ A semantics. |
SDEKernel |
class |
Abstract base class for state-space kernel representations. |
SDEParams |
class |
Continuous-time SDE parameters for a linear SDE. |
SLQLogdet |
class |
Stochastic log-determinant via Lanczos quadrature (SLQ). |
SLRResult |
class |
Linear-Gaussian surrogate produced by statistical linear regression. |
SVDLowRankUpdate |
class |
Deprecated subclass of LowRankUpdate with orthonormal=True. |
ScaledOperator |
function |
Lazy scalar multiply (c A) v = c (A v). |
SoftmaxLikelihood |
class |
Softmax (categorical) likelihood for multi-class classification. |
StudentTLikelihood |
class |
Student-t likelihood for robust regression. |
SumKronecker |
class |
Deprecated alias for SumOfKroneckers. |
SumKroneckerSqrt |
class |
Lazy Lanczos square-root operator for SumOfKroneckers covariances. |
SumOfKroneckers |
class |
Sum of Kronecker products Σ_k A_k \otimes B_k. |
SumOperator |
function |
Lazy sum (A + B + …) v = A v + B v + …. |
SumSDE |
class |
Sum of SDE kernels via block-diagonal composition. |
TaylorIntegrator |
class |
1st or 2nd order Taylor expansion for uncertainty propagation. |
Toeplitz |
class |
Symmetric Toeplitz matrix from its first column. |
ToeplitzCholesky |
class |
Circulant-embedding sample factor for a symmetric positive Toeplitz matrix. |
UDLDecomposition |
class |
Banded factorisation \(\Lambda = U \tilde{D} U^{\top}\) of a BlockTriDiag. |
UnscentedIntegrator |
class |
Unscented transform: deterministic sigma points. |
UpperBlockTriDiag |
class |
Upper triangular block-bidiagonal (transpose of LowerBlockTriDiag). |
add_jitter |
function |
Add diagonal jitter for numerical stability: A + eps * I. |
as_diagonalised |
function |
Express operator as one DiagonalisedOperator, if possible. |
as_linear_operator |
function |
Wrap a raw matvec callable as a tagged lineax operator. |
base_conditional |
function |
Gaussian conditional distribution via Schur complement. |
batched_kernel_matvec |
function |
Compute K(X, Z) @ v in memory-efficient batches. |
batched_kernel_rmatvec |
function |
Compute K(X, Z)^T @ u in memory-efficient batches. |
block_diagonal_tag |
constant |
block_diagonal_tag |
block_tridiagonal_tag |
constant |
block_tridiagonal_tag |
blr_diag_update |
function |
Diagonal natural parameter BLR update step. |
blr_full_update |
function |
Full-rank natural parameter BLR update step. |
build_prediction_cache |
function |
Solve A alpha = y and cache the result. |
cavity_distribution |
function |
Compute EP cavity distribution by removing a site. |
cavity_from_marginal |
function |
Compute cavity distribution by removing a site from the marginal. |
center_kernel |
function |
Center a kernel matrix: H K H. |
centering_operator |
function |
Centering matrix H = I - (1/n) \mathbf{1}\mathbf{1}^T. |
cholesky |
function |
Compute Cholesky factor L such that A = L L^T. |
cholesky_logdet |
function |
Compute log|A| from Cholesky factor L where A = L Lᵀ. |
circulant_from_symbol |
function |
(Block-)circulant operator from its DFT symbol. |
collapsed_elbo |
function |
Collapsed ELBO (Titsias bound) for sparse GP regression. |
compute_psi_statistics |
function |
Compute Ψ statistics, dispatching to analytical or numerical. |
conditional |
function |
Compute p(x_A \| x_B = b) from a joint Gaussian p(x_A, x_B). |
conditional_interpolate |
function |
Interpolated marginal at time t given posteriors at t^- and t^+. |
conditional_variance |
function |
Predictive variance: Schur complement diagonal plus optional variational correction. |
cost_expectation |
function |
Compute E[Cost(f(x), target)] where x ~ N(mu, Sigma). |
cov_transform |
function |
Covariance propagation through a linear map: J @ Sigma @ J^T. |
create_grid |
function |
Create a regular grid from per-dimension sizes and bounds. |
cubature_points |
function |
Spherical-radial cubature points and weights. |
cubic_interpolation_weights |
function |
Compute cubic interpolation indices and weights for SKI. |
cvi_update_sites |
function |
Natural gradient update for CVI sites. |
damped_natural_update |
function |
Damped update in natural parameter space. |
dare |
function |
Discrete Algebraic Riccati Equation solver. |
diag |
function |
Extract the diagonal of an operator as a 1D array. |
diag_conditional_variance |
function |
Diagonal of Schur complement: diag(K_XX - A K_ZX). |
diag_inv |
function |
Compute the diagonal of the inverse of a linear operator. |
diagonal_tag |
constant |
Re-exported from lineax. |
discrepancy_step_size |
function |
Adaptive EKI tempering step: the data misfit controller. |
discrete_lyapunov_solve |
function |
Solve the discrete Lyapunov equation P - G P G^T = Q. |
discretise_mfd |
function |
Discretise a linear SDE by matrix-fraction decomposition. |
discretise_mfd_sequence |
function |
Vectorised gaussx.discretise_mfd over a vector of time steps. |
dist_kl_divergence |
function |
KL divergence KL(p \|\| q) between two multivariate normals. |
eig |
function |
Compute eigenvalues and eigenvectors. |
eigenpro_correction |
function |
Compute the EigenPro eigenspace correction for one SGD mini-batch. |
eigenpro_preconditioner |
function |
Build an EigenPro spectral preconditioner from a kernel operator. |
eigenpro_step_size |
function |
Return the EigenPro preconditioned step size for a batch size. |
eigvals |
function |
Compute eigenvalues only. |
eki_step |
function |
One ensemble Kalman inversion (EKI) update. |
elbo |
function |
Evidence lower bound (ELBO). |
enkf_analysis |
function |
Stochastic (perturbed-observation) ensemble Kalman analysis step. |
ensemble_covariance |
function |
Empirical covariance from an ensemble as a low-rank operator. |
ensemble_cross_covariance |
function |
Cross-covariance between two ensemble sets. |
ensemble_kalman_gain |
function |
Kalman gain from an ensemble and its image in observation space. |
ep_tilted_moments |
function |
Compute tilted distribution moments via Gauss-Hermite quadrature. |
estimate_spectral_bounds |
function |
Estimate \((\lambda_{\min}, \lambda_{\max})\) of a symmetric operator. |
etkf_transform |
function |
Ensemble Transform Kalman Filter (ETKF) analysis weights. |
euclidean_distance |
function |
Pairwise Euclidean distances \|\|a_i - b_j\|\|. |
expectation_to_meanvar |
function |
Convert expectation parameters to mean/variance (Cholesky). |
expectation_to_natural |
function |
Convert expectation parameters to natural parameters. |
expectations_to_ssm |
function |
Convert expectation parameters back to SSM marginals. |
expected_log_likelihood |
function |
Unified expected log-likelihood with analytical dispatch. |
falkon_preconditioner |
function |
Build the Falkon preconditioner from the inducing-point kernel matrix. |
falkon_predict |
function |
Predict with Nyström KRR weights: \(f(x_*) = \sum_j \alpha_j\, k(x_*, z_j)\). |
falkon_solve |
function |
Nyström kernel ridge regression weights by Falkon's preconditioned CG. |
fifth_order_cubature_points |
function |
Fifth-order fully symmetric cubature points and weights. |
fisher_info |
function |
Fisher information matrix F(eta) = nabla^2 A(eta). |
frobenius_norm |
function |
Compute the Frobenius norm \|\|A\|\|_F with structural dispatch. |
gaspari_cohn |
function |
Gaspari-Cohn (1999) fifth-order compactly-supported taper. |
gauss_hermite_points |
function |
Gauss-Hermite quadrature points and weights. |
gauss_kl |
function |
KL divergence KL[q(u) \|\| p(u)] between Gaussian distributions. |
gauss_newton_precision |
function |
Gauss-Newton precision matrix J^T J. |
gaussian_entropy |
function |
Entropy of a multivariate normal N(mu, Sigma). |
gaussian_expected_log_lik |
function |
Expected log-likelihood E_q[log N(y \| f, R)]. |
gaussian_log_prob |
function |
Multivariate normal log-probability. |
ggn_diagonal |
function |
Generalized Gauss-Newton diagonal approximation. |
gradient_expectation |
function |
Compute E[nabla f(x)] via Stein's lemma. |
grid_coupling_indices |
function |
Flat indices of masked-out grid cells adjacent to the masked-in region. |
grid_data |
function |
Expand a grid to the full Cartesian product of grid points. |
haversine_distance |
function |
Pairwise great-circle (haversine) distances on a sphere. |
hsic |
function |
Biased HSIC estimator. |
hutchinson_hessian_diag |
function |
Stochastic Hessian diagonal via Hutchinson with Rademacher probes. |
implicit_cross_kernel |
function |
Create a matrix-free rectangular cross-kernel operator. |
infinite_horizon_filter |
function |
Infinite-horizon Kalman filter with fixed steady-state gain. |
infinite_horizon_smoother |
function |
Infinite-horizon RTS smoother with fixed steady-state gain. |
inflate_multiplicative |
function |
Multiplicative ensemble inflation about the mean. |
inflate_rtpp |
function |
Relaxation to prior perturbations (RTPP; Zhang et al. 2004). |
inflate_rtps |
function |
Relaxation to prior spread (RTPS; Whitaker & Hamill 2012). |
inv |
function |
Return a lazy inverse operator A^{-1}. |
inv_quad_logdet |
function |
Joint inverse-quadratic and log-determinant via shared BBMM CG. |
is_block_diagonal |
function |
Check whether operator carries the block-diagonal tag. |
is_block_tridiagonal |
function |
Check whether operator carries the block-tridiagonal tag. |
is_diagonal |
function |
Re-exported from lineax. |
is_kronecker |
function |
Check whether operator carries the Kronecker tag. |
is_kronecker_sum |
function |
Check whether operator carries the Kronecker sum tag. |
is_low_rank |
function |
Check whether operator carries the low-rank tag. |
is_lower_triangular |
function |
Re-exported from lineax. |
is_negative_semidefinite |
function |
Re-exported from lineax. |
is_positive_semidefinite |
function |
Re-exported from lineax. |
is_symmetric |
function |
Re-exported from lineax. |
is_upper_triangular |
function |
Re-exported from lineax. |
joseph_update |
function |
Numerically stable Joseph-form covariance update. |
kalman_filter |
function |
Kalman filter forward pass via jax.lax.scan. |
kalman_gain |
function |
Compute Kalman gain K = P @ H^T @ (H @ P @ H^T + R)^{-1}. |
kernel_expectations |
function |
Compute kernel expectations Psi_0, Psi_1, Psi_2 for uncertain inputs. |
kl_divergence |
function |
KL divergence KL(q \|\| p) via the Bregman-divergence form on natural parameters. |
kl_standard_normal |
function |
KL divergence KL(N(m, S) \|\| N(0, I)). |
kronecker_mll |
function |
Exact marginal log-likelihood for a Kronecker-structured GP. |
kronecker_posterior_predictive |
function |
Posterior mean and variance for a Kronecker GP at test points. |
kronecker_sum_sample |
function |
Sample from 𝒩(0, A ⊕ B) using per-factor eigendecompositions. |
kronecker_sum_solve |
function |
Solve a shifted Kronecker-sum system in tensor (unvectorized) form. |
kronecker_sum_tag |
constant |
kronecker_sum_tag |
kronecker_tag |
constant |
kronecker_tag |
leave_one_out_cv |
function |
LOO-CV via the bordered-system identity. |
linear_solve |
function |
Solve A x = b through the unified front door. |
localization_matrix |
function |
Pairwise Gaspari-Cohn taper rho(dist(a_i, b_j); c). |
localized_kalman_gain |
function |
Ensemble Kalman gain with Hadamard (Schur-product) localization. |
log_likelihood_expectation |
function |
Compute E[log p(y_{obs} \| f(x))] where x ~ N(mu, Sigma). |
log_marginal_likelihood |
function |
GP log marginal likelihood. |
log_partition |
function |
Log-partition function A(eta). |
logdet |
function |
Compute log |det(A)| with structural dispatch. |
love_cache |
function |
Precompute Lanczos factorization of K^{-1} for fast variance. |
love_variance |
function |
Fast predictive variance using a LOVE cache. |
low_rank_plus_diag |
function |
Construct diag(diag) + U diag(d) Vᵀ. |
low_rank_plus_identity |
function |
Construct scale * I + U diag(d) Vᵀ. |
low_rank_tag |
constant |
low_rank_tag |
lower_triangular_tag |
constant |
Re-exported from lineax. |
masked_moment_inputs |
function |
Make masked observation channels inert in a moment-matched update. |
matheron_update |
function |
Posterior samples via Matheron's-rule correction. |
mean_cov_to_natural |
function |
Convert mean/covariance to natural parameters (operator form). |
mean_expectation |
function |
Compute E[f(x)] where x ~ N(mu, Sigma). |
meanfield_kalman_filter |
function |
Mean-field Kalman filter over L = D // block_size state blocks. |
meanfield_rts_smoother |
function |
Mean-field RTS smoother over L = D // block_size state blocks. |
meanvar_to_expectation |
function |
Convert mean/variance (Cholesky) to expectation parameters. |
meanvar_to_natural |
function |
Convert mean/variance (Cholesky) to natural parameters. |
mmd_squared |
function |
Biased squared Maximum Mean Discrepancy. |
moment_match |
function |
EP moment matching: tilted log-normaliser and its first two derivatives. |
moment_transform |
function |
Moment triple of (x, fn(x)) for x Gaussian. |
natural_to_expectation |
function |
Convert natural parameters to expectation parameters. |
natural_to_mean_cov |
function |
Convert natural parameters to mean/covariance (operator form). |
natural_to_meanvar |
function |
Convert natural parameters to mean/variance (Cholesky). |
naturals_to_ssm |
function |
Convert natural parameters back to SSM parameters. |
negative_semidefinite_tag |
constant |
Re-exported from lineax. |
newton_update |
function |
Convert a Newton step to natural pseudo-likelihood parameters. |
nonlinear_kalman_filter |
function |
Moment-matched nonlinear Kalman filter. |
nonlinear_kalman_predict |
function |
One moment-matched predict step. |
nonlinear_kalman_update |
function |
One moment-matched update step. |
nonlinear_rts_smoother |
function |
Moment-matched nonlinear Rauch-Tung-Striebel smoother. |
nonlinear_rts_step |
function |
One moment-matched RTS backward step. |
nystrom_operator |
function |
Nystrom low-rank kernel approximation. |
oilmm_back_project |
function |
Back-project latent GP predictions to the observation space. |
oilmm_project |
function |
Project multi-output data to independent latent GPs via OILMM. |
pairwise_marginals |
function |
Joint p(x_k, x_{k+1}) for each consecutive pair. |
parallel_kalman_filter |
function |
Parallel Kalman filter via jax.lax.associative_scan. |
parallel_rts_smoother |
function |
Parallel RTS smoother via reverse jax.lax.associative_scan. |
positive_semidefinite_tag |
constant |
Re-exported from lineax. |
predict_mean |
function |
Predictive mean: mu* = K_*f @ alpha. |
predict_variance |
function |
Predictive variance: sigma^2* = k_** - diag(K_*f K_y^{-1} K_f*). |
process_noise_covariance |
function |
Compute process noise from stationary covariance. |
project |
function |
Compute A_X = K_XZ @ K_ZZ^{-1} via Cholesky solve. |
quadratic_form |
function |
Compute x^T A^{-1} x via a single solve. |
rff_operator |
function |
Random Fourier Features kernel approximation. |
riemannian_psd_correction |
function |
Riemannian gradient correction for PSD precision updates. |
root_decomposition |
function |
Compute a tall factor R such that R Rᵀ ≈ A. |
root_inv_decomposition |
function |
Compute a tall factor R⁻ such that R⁻ (R⁻)ᵀ ≈ A⁻¹. |
rts_smoother |
function |
Rauch-Tung-Striebel backward smoother. |
safe_cholesky |
function |
Cholesky decomposition with adaptive jitter for near-singular matrices. |
sample_joint_conditional |
function |
Joint draws of a partitioned Gaussian, and conditional draws from them. |
sample_mvn |
function |
Draw samples from \(\mathcal{N}(\mu, K)\), dispatching on \(K\)'s structure. |
sandwich |
function |
Return A @ P @ A.T exploiting compatible operator structure. |
schur_complement |
function |
Schur complement: K_XX - K_XZ @ K_ZZ^{-1} @ K_ZX. |
sde_autocovariance |
function |
Compute the stationary autocovariance of an SDE kernel. |
sigma_points |
function |
Unscented transform sigma points and weights. |
site_mean_var_from_natural |
function |
Convert per-site natural parameters to mean/variance. |
site_natural_from_tilted |
function |
Compute site natural parameters from tilted and cavity moments. |
sites_to_precision |
function |
Convert Gaussian sites to a block-tridiagonal precision. |
solve |
function |
Solve A x = b with structural dispatch. |
solve_columns |
function |
Solve A X = B column-by-column via vmap. |
solve_matrix |
function |
Solve A X = B with a single factorization on the matrix RHS. |
solve_rows |
function |
Solve A x = bᵢ for each row bᵢ of a matrix via vmap. |
solve_tridiagonal |
function |
Solve a tridiagonal system A x = d. |
solve_tridiagonal_batched |
function |
Solve independent tridiagonal systems over leading batch dimensions. |
spingp_log_likelihood |
function |
Log marginal likelihood via sparse inverse GP formulation. |
spingp_posterior |
function |
Posterior mean and precision via SpInGP. |
sqrt |
function |
Compute matrix square root S such that S @ S = A. |
sqrt_inv_matmul |
function |
Compute \(A^{-1/2} b\) via contour-integral quadrature. |
sqrt_matmul |
function |
Compute \(A^{1/2} b\) via contour-integral quadrature. |
ssm_to_expectations |
function |
Convert SSM marginals to expectation parameters. |
ssm_to_naturals |
function |
Convert SSM parameters to natural parameters. |
stable_rbf_kernel |
function |
RBF (squared exponential) kernel with mixed-precision stability. |
stable_squared_distances |
function |
Squared Euclidean distances with mixed-precision stability. |
statistical_linear_regression |
function |
Moment-match a nonlinear conditional to a linear-Gaussian surrogate. |
submatrix |
function |
Extract A[row_idx, col_idx] without forming the full matrix. |
sufficient_stats |
function |
Compute sufficient statistics T(x) = [x, x x^T]. |
sumkronecker_sample |
function |
Sample from 𝒩(0, op) with matrix-free Lanczos square roots. |
svd |
function |
Compute the singular value decomposition A = U diag(s) V^T. |
svd_low_rank_plus_diag |
function |
Construct diag(diag) + U diag(S) Vᵀ from a truncated SVD. |
svgp_variance_adjustment |
function |
Compute the SVGP variance adjustment operator. |
symmetric_tag |
constant |
Re-exported from lineax. |
symmetrize |
function |
Symmetrize the trailing two axes: 0.5 * (X + X^T). |
tikhonov_augment |
function |
Observation augmentation for Tikhonov-regularised EKI (TEKI). |
to_expectation |
function |
Convert natural to expectation parameters. |
to_natural |
function |
Convert expectation to natural parameters. |
toeplitz_sample |
function |
Sample from 𝒩(0, Toeplitz(column)) via FFT circulant embedding. |
trace |
function |
Compute the trace of an operator. |
trace_and_diag |
function |
Jointly estimate the trace and diagonal from one probe pass. |
trace_correction |
function |
Trace term in Titsias collapsed ELBO. |
trace_product |
function |
Trace of a matrix product: tr(A @ B) with structural dispatch. |
tridiagonal_tag |
constant |
Re-exported from lineax. |
udl_decomposition |
function |
Factorise a symmetric block-tridiagonal precision as \(U \tilde{D} U^{\top}\). |
udl_from_ssm_params |
function |
Build the UDLDecomposition of a Gauss-Markov chain's precision. |
udl_to_ssm_params |
function |
Read the equivalent Gauss-Markov chain off a UDLDecomposition. |
uncertain_bgplvm_predict |
function |
Multi-output uncertain GP prediction for BGPLVM. |
uncertain_gp_predict |
function |
Predictive mean and variance for GP with uncertain inputs. |
uncertain_gp_predict_mc |
function |
Monte Carlo GP prediction with uncertain inputs. |
uncertain_svgp_predict |
function |
Predictive mean and variance for SVGP with uncertain inputs. |
uncertain_vgp_predict |
function |
Predictive mean and variance for dense VGP with uncertain inputs. |
unit_diagonal_tag |
constant |
Re-exported from lineax. |
unwhiten |
function |
Unwhiten variational mean: m = L @ m_tilde. |
unwhiten_covariance |
function |
Unwhiten variational covariance: S = L S̃ Lᵀ. |
upper_triangular_tag |
constant |
Re-exported from lineax. |
variational_elbo_gaussian |
function |
Titsias collapsed ELBO for Gaussian likelihoods. |
variational_elbo_mc |
function |
Monte Carlo ELBO for non-conjugate likelihoods. |
whiten_covariance |
function |
Unwhiten variational covariance: S = L S̃ Lᵀ. |
whitened_svgp_predict |
function |
Whitened SVGP prediction: mean and variance at test points. |
woodbury_solve |
function |
Standalone Woodbury identity solve: (L + U diag(D) U^T)^{-1} b. |