Capability index

Every public name in the four pyrox packages (245 of them), each listed once at its public home and grouped by the module that defines it, with the first sentence of its docstring; then the public API of gaussx, kernellib and geonnax, which the packages build on. Search this page before writing a helper: if what you need is here, compose it; if it is almost here, extend it where it lives. Regenerate with make capabilities after changing a public API (packages/pyrox/tests/test_capabilities.py checks it is current).

pyrox

pyrox._core

Declarative wrappers for NumPyro param and sample sites. (pyrox._core.descriptors)

Name	Kind	What it does
PyroxParam	class	Lightweight metadata container for a parameter site.
PyroxSample	class	Lightweight metadata container for a random sample site.

Parameterized base class for GP-style modules. (pyrox._core.parameterized)

Name	Kind	What it does
Parameterized	class	Shared base for modules with priors, constraints, and mode switching.

PyroxModule and pyrox_method — Equinox-to-NumPyro bridge primitives. (pyrox._core.pyrox_module)

Name	Kind	What it does
PyroxModule	class	Equinox module with NumPyro site registration and per-call caching.
pyrox_method	function	Wrap a method so its body runs inside the module's per-call context.

pyrox.inference

Ensemble-of-MAP / ensemble-of-VI primitives. (pyrox.inference._ensemble)

Name	Kind	What it does
EnsembleMAP	class	NumPyro-like ensemble MAP/MLE runner.
EnsembleResult	class	Output of a full EnsembleMAP.run / EnsembleVI.run.
EnsembleState	class	Stacked state for an ensemble of optimizer runs.
EnsembleVI	class	NumPyro-like ensemble variational-inference runner.
ensemble_init	function	Initialize an ensemble of (params, opt_state) by vmap over keys.
ensemble_loss	function	Build the filter_value_and_grad loss from a log-joint.
ensemble_map	function	One-shot wrapper around EnsembleMAP.
ensemble_predict	function	Vmap predict_fn over the leading ensemble axis of params.
ensemble_step	function	Perform one ensemble update step on a batch.
ensemble_vi	function	One-shot wrapper around EnsembleVI.

Post-hoc last-layer Laplace approximation with a GGN Hessian. (pyrox.inference._laplace)

Name	Kind	What it does
LastLayerLaplace	class	Gaussian posterior over a selected parameter group of a fitted model.
last_layer_laplace	function	Fit a post-hoc GGN Laplace posterior over a parameter group.

Per-parameter-group optimizer helper. (pyrox.inference._param_groups)

Name	Kind	What it does
param_group_optimizer	function	Apply a different optimizer to each labelled parameter group.

pyrox-gp

pyrox_gp

Sparse SVGP variational guide families — concrete Guide types. (pyrox_gp._guides)

Name	Kind	What it does
DeltaGuide	class	Point-mass (MAP-style) variational posterior over inducing values.
FullRankGuide	class	Full-rank Gaussian variational posterior over inducing values u.
MeanFieldGuide	class	Diagonal Gaussian variational posterior over inducing values u.
NaturalGuide	class	Natural-parameter Gaussian variational posterior over inducing values.
WhitenedGuide	class	Whitened-coordinate Gaussian variational posterior over inducing values.

Inter-domain inducing-feature families for sparse GPs. (pyrox_gp._inducing)

Name	Kind	What it does
DecoupledInducingFeatures	class	Decoupled mean / covariance inducing-feature bases (Cheng & Boots 2017).
FourierInducingFeatures	class	VFF — Variational Fourier inducing features on \([-L, L]^D\).
InducingFeatures	class	Protocol for inter-domain inducing features.
LaplacianInducingFeatures	class	Inducing features from low-frequency graph Laplacian eigenvectors.
SlepianInducingFeatures	class	Region-localized Slepian inducing features on \(S^2\).
SphericalHarmonicInducingFeatures	class	VISH — inducing harmonics on \(S^2\) (Dutordoir et al. 2020).
funk_hecke_coefficients	function	Funk-Hecke coefficients of a zonal kernel on \(S^2\).

Choosing inducing inputs from the data (P3). (pyrox_gp._inducing_init)

Name	Kind	What it does
init_inducing	function	Inducing inputs (M, D): rows of X picked by kernellib.select_landmarks.

Sparse variational GP inference entry points. (pyrox_gp._inference)

Name	Kind	What it does
ConjugateVI	class	Natural-gradient / CVI update for sparse variational GPs.
svgp_elbo	function	Structured SVGP ELBO as a differentiable scalar.
svgp_factor	function	Register the structured SVGP ELBO as a NumPyro factor site.

Advanced GP inference strategies for non-conjugate likelihoods. (pyrox_gp._inference_nongauss)

Name	Kind	What it does
ExpectationPropagation	class	Parallel expectation propagation (Minka 2001) with damping.
GaussNewtonInference	class	Gauss-Newton inference: Newton loop with PSD-projected curvature.
LaplaceInference	class	Laplace approximation via Newton iteration on the log posterior.
NonGaussConditionedGP	class	GP conditioned on a non-Gaussian likelihood via an advanced strategy.
PosteriorLinearization	class	Iterated statistical-linearization site updates.
QuasiNewtonInference	class	MAP optimization via L-BFGS, Laplace covariance at convergence.

Site-based non-Gaussian inference for MarkovGPPrior. (pyrox_gp._inference_nongauss_markov)

Name	Kind	What it does
ExpectationPropagationMarkov	class	Parallel Expectation Propagation for MarkovGPPrior.
GaussNewtonMarkovInference	class	Gauss-Newton (Markov) inference: Newton with a strict PSD floor.
LaplaceMarkovInference	class	Laplace approximation for MarkovGPPrior.
NonGaussConditionedMarkovGP	class	Markov GP conditioned on a non-Gaussian likelihood via a site-based strategy.
PosteriorLinearizationMarkov	class	Posterior linearization for MarkovGPPrior (Markov IPLF).

Concrete kernel classes — Parameterized wrappers over the math primitives. (pyrox_gp._kernels)

Name	Kind	What it does
Constant	class	Constant kernel k(x, x') = sigma^2.
Cosine	class	Cosine kernel sigma^2 cos(2 pi \|\|x - x'\|\| / period).
Linear	class	Linear kernel sigma^2 x^T x' + bias.
Matern	class	Matern kernel with nu in {0.5, 1.5, 2.5}.
Periodic	class	Periodic (MacKay) kernel.
Polynomial	class	Polynomial kernel sigma^2 (x^T x' + bias)^degree.
RBF	class	Radial basis function (squared exponential) kernel.
RationalQuadratic	class	Rational quadratic kernel.
White	class	White-noise kernel sigma^2 delta(x, x').

Collapsed latent-factor regression — the decoder is marginalized, not fixed. (pyrox_gp._latent_factor)

Name	Kind	What it does
collapsed_lfr_log_prob	function	Log-likelihood with the linear decoder marginalized analytically.
decoder_posterior	function	Closed-form matrix-normal posterior over the decoder.
lfr_predictive_moments	function	Exact moments of the product of two independent Gaussians.
warp_to_base	function	Map observations into the space where the factor model is linear.
warped_decoder_posterior	function	Matrix-normal decoder posterior, read in the warped space.
warped_lfr_log_prob	function	Collapsed latent-factor log-likelihood on warped observations.

Model surface for collapsed latent-factor regression. (pyrox_gp._latent_factor_models)

Name	Kind	What it does
ConditionedLatentFactorGP	class	Latent-factor posterior — MAP latents plus the decoder posterior.
LatentFactorGPPrior	class	GP latent factors with an analytically marginalized linear decoder.
latent_total_correlation	function	Total correlation of the fitted latent factors.
lfr_factor	function	Register the collapsed latent-factor log-likelihood with NumPyro.
lfr_model	function	MAP-over-latents collapsed latent-factor regression model.

Data-driven initialisation of GPLVM-style latent variables (P1). (pyrox_gp._latent_init)

Name	Kind	What it does
latent_init	function	Initial latents (N, n_latent) for a GPLVM or latent-factor model.

Observation models for GP inference. (pyrox_gp._likelihoods)

Name	Kind	What it does
BernoulliLikelihood	class	Binary classification likelihood with logit link.
DistLikelihood	class	Generic likelihood wrapping any numpyro.distributions.Distribution.
GaussianLikelihood	class	Gaussian observation model \(p(y \mid f) = N(y \mid f, \sigma^2)\).
HeteroscedasticGaussianLikelihood	class	Gaussian regression with input-dependent noise.
PoissonLikelihood	class	Count likelihood with log-link.
SoftmaxLikelihood	class	Multi-class classification likelihood with softmax link.
StudentTLikelihood	class	Heavy-tailed regression: p(y \| f) = StudentT(nu, f, sigma).

Kalman-based Markov GP prior on top of the SDE-kernel surface. (pyrox_gp._markov)

Name	Kind	What it does
ConditionedMarkovGP	class	Markov GP conditioned on Gaussian-likelihood observations.
MarkovGPPrior	class	Linear-time temporal GP prior over a sorted 1-D grid.
markov_gp_factor	function	Register the collapsed Markov-GP marginal log-likelihood with NumPyro.
markov_gp_sample	function	Sample a latent function f at the prior's training times.

Normalizing Kalman Filter prior on top of the gaussx state-space surface. (pyrox_gp._markov_flow)

Name	Kind	What it does
NormalizingKalmanPrior	class	Normalizing Kalman Filter prior over a multivariate time grid.
normalizing_kalman_factor	function	Register the NKF marginal log-likelihood with NumPyro.

Model-facing GP entry points — GPPrior, ConditionedGP, and NumPyro hooks. (pyrox_gp._models)

Name	Kind	What it does
ConditionedGP	class	GP conditioned on Gaussian-likelihood training observations.
GPPrior	class	Finite-dimensional GP prior over a fixed training input set.
GradientPosterior	class	Gaussian posterior over \(\nabla f\) at \(M\) test points (point-major).
gp_factor	function	Register the collapsed GP log marginal likelihood with NumPyro.
gp_sample	function	Sample a latent function f at the prior's training inputs.

Multi-output Gaussian-process kernels and inducing structures. (pyrox_gp._multi_output)

Name	Kind	What it does
ICMKernel	class	Intrinsic coregionalization model with one shared latent kernel.
LMCKernel	class	Linear model of coregionalization for vector-valued GPs.
MultiOutputInducingVariables	class	Shared inducing-point structure for LMC-style sparse workflows.
OILMMKernel	class	Orthogonal instantaneous linear mixing model.
SharedInducingPoints	class	Shared inducing inputs for multi-output latent GP collections.

Multi-output GP model entry points — exact and sparse variational. (pyrox_gp._multi_output_models)

Name	Kind	What it does
ConditionedOILMMGP	class	OILMM posterior — Q conditioned scalar GPs plus the mixing.
MultiOutputConditionedGP	class	Multi-output GP conditioned on Gaussian-likelihood observations.
MultiOutputGPPrior	class	Finite-dimensional multi-output GP prior over fixed training inputs.
MultiOutputSparseGPPrior	class	Sparse multi-output GP prior over shared inducing inputs.
OILMMGPPrior	class	Orthogonal-mixing multi-output GP prior with a projected exact posterior.
mo_gp_factor	function	Register the collapsed multi-output GP marginal likelihood with NumPyro.
mo_svgp_elbo	function	Structured multi-output SVGP ELBO as a differentiable scalar.
mo_svgp_factor	function	Register the multi-output SVGP ELBO as a NumPyro factor site.

Pathwise GP posterior samplers via Matheron's rule. (pyrox_gp._pathwise)

Name	Kind	What it does
DecoupledPathwiseSampler	class	Sparse/decoupled pathwise sampler with RFF prior + inducing update.
PathwiseFunction	class	Callable posterior function draw(s) produced by a pathwise sampler.
PathwiseSampler	class	Exact-GP pathwise posterior sampler using Matheron's rule.

A gaussx solver strategy for exact GPs at large n (P4). (pyrox_gp._preconditioned)

Name	Kind	What it does
preconditioned_cg_solver	function	A gaussx solver strategy for exact GPs: preconditioned CG plus SLQ.

Layer 1 — abstract protocol classes for GP components. (pyrox_gp._protocols)

Name	Kind	What it does
Guide	class	Abstract base for variational posterior families.
Likelihood	class	Abstract base for observation models.

Sparse GP prior — inducing-input variational foundation. (pyrox_gp._sparse)

Name	Kind	What it does
SparseGPPrior	class	GP prior parameterized over inducing inputs Z.

Sparse variational Markov-GP — inducing time points on top of an SDE prior. (pyrox_gp._sparse_markov)

Name	Kind	What it does
SparseConditionedMarkovGP	class	Sparse Markov GP fitted to a guide.
SparseMarkovGPPrior	class	Sparse variational GP prior over an SDE kernel and inducing time grid.
sparse_markov_elbo	function	Sparse variational ELBO for SparseMarkovGPPrior.
sparse_markov_factor	function	Register sparse_markov_elbo as a NumPyro factor site.

Transformed Gaussian processes: a GP composed with a monotone warp. (pyrox_gp._warped)

Name	Kind	What it does
WarpedGaussianLikelihood	class	\(p(y \mid f) = N(y \mid G(f), \sigma^2)\) -- the transformed-GP model.
warped_predictive_moments	function	Moments of \(y = G(f) + \epsilon\) under \(q(f)\), by quadrature.

Re-exported from upstream

Name	Kind	What it does
ConstantSDE	class	Re-exported from gaussx.
CosineSDE	class	Re-exported from gaussx.
IntegratedWienerSDE	class	Re-exported from gaussx.
Kernel	class	Re-exported from kernellib.
MaternSDE	class	Re-exported from gaussx.
PeriodicSDE	class	Re-exported from gaussx.
ProductSDE	class	Re-exported from gaussx.
QuasiPeriodicSDE	class	Re-exported from gaussx.
SDEKernel	class	Re-exported from gaussx.
SDEParams	class	Re-exported from gaussx.
SumSDE	class	Re-exported from gaussx.

pyrox-nn

pyrox_nn

Bayesian Neural Field (BNF) layers — port of Google's bayesnf. (pyrox_nn._bnf)

Name	Kind	What it does
BayesianNeuralField	class	The full Bayesian Neural Field architecture, wrapping geonnax.NeuralFieldMLP.
FourierFeatures	class	Per-input dyadic-frequency Fourier basis.
InteractionFeatures	class	Element-wise products on selected pairs of input columns.
SeasonalFeatures	class	Period-and-harmonic cos/sin basis on a scalar time axis.
Standardization	class	Apply a fixed-coefficient affine standardization.

Bayesian conditioning primitives for pyrox_nn. (pyrox_nn._conditioning)

Name	Kind	What it does
BayesianAffineModulation	class	geonnax.AffineModulation with Normal priors on the FiLM generator.
BayesianConcatConditioner	class	geonnax.ConcatConditioner with Normal priors on the projection.
BayesianHyperLinear	class	geonnax.HyperLinear with Normal priors on the generator only.
ConditionedRFFNet	class	Conditional analogue of pyrox_nn.RandomKitchenSinks.
HyperFourierFeatures	class	Random Fourier features with (W, b, log_lengthscale) from a parameter net.

Bayesian / uncertainty-aware dense layers. (pyrox_nn._dense)

Name	Kind	What it does
DenseDVI	class	Deterministic Variational Inference dense layer (Wu et al., 2018).
DenseFlipout	class	Bayesian dense layer with Flipout sign-flip structure.
DenseHierarchical	class	Hierarchical Bayesian dense layer with multiplicative shrinkage.
DenseNCP	class	Noise Contrastive Prior dense layer (Hafner et al., 2019).
DenseReparameterization	class	Bayesian dense layer via the reparameterization trick.
DenseVariational	class	Dense layer with a user-supplied prior factory.
DenseVariationalDropout	class	Sparse variational dropout dense layer.
NCPNormalOutput	class	Output-side Noise Contrastive Prior layer (Hafner et al., 2018).

Ensemble-style layers for pyrox_nn. (pyrox_nn._ensemble)

Name	Kind	What it does
DenseRank1	class	Rank-1 ensemble dense layer.
LayerNormEnsemble	class	Per-ensemble-member LayerNorm.
MultiHeadAttentionBE	class	Multi-head attention with BatchEnsemble rank-1 projections.

Pure-JAX feature helpers and Bayesian RFF feature layers. (pyrox_nn._features)

Name	Kind	What it does
ArcCosineFourierFeatures	class	Random features for the arc-cosine kernel (Cho & Saul, 2009).
HSGPFeatures	class	Hilbert-Space Gaussian Process feature layer (Riutort-Mayol et al., 2023).
LaplaceCosineFeatures	class	Cosine-bias variant of random Fourier features for the Laplace kernel.
LaplaceFourierFeatures	class	SSGP-style RFF layer with Laplace (Matern-1/2) spectral density.
MaternCosineFeatures	class	Cosine-bias variant of random Fourier features for the Matern kernel.
MaternFourierFeatures	class	SSGP-style RFF layer with Matern spectral density.
RBFCosineFeatures	class	Cosine-bias variant of random Fourier features for the RBF kernel.
RBFFourierFeatures	class	SSGP-style RFF layer with RBF spectral density.
RandomKitchenSinks	class	Random Kitchen Sinks: RFF + a learned linear head.
SphericalHarmonicFeatures	class	Weight-space GP on the sphere \(S^2\) from real spherical harmonics.
VariationalFourierFeatures	class	VSSGP — RFF with a learnable variational posterior over frequencies.

Heteroscedastic output layers (Collier et al., 2021). (pyrox_nn._heteroscedastic)

Name	Kind	What it does
GaussianRegressionHead	class	Bayesian heteroscedastic regression head (mean plus log-scale).
MCSigmoidDenseFA	class	Heteroscedastic multi-label output layer (FA noise + sigmoid).
MCSoftmaxDenseFA	class	Heteroscedastic multi-class output layer (FA noise + softmax).

Bayesian Multiplicative Filter Networks. (pyrox_nn._mfn)

Name	Kind	What it does
BayesianFourierNet	class	FourierNet with Bayesian priors on all filter and linear weights.
BayesianGaborNet	class	GaborNet with Bayesian priors on all filter and linear weights.

Bayesian SIREN — sinusoidal network with regime-scaled Normal priors. (pyrox_nn._siren)

Name	Kind	What it does
BayesianSIREN	class	SIREN with regime-scaled Normal priors on all layer weights.

Bayesian Slepian positional encoder. (pyrox_nn._slepian)

Name	Kind	What it does
BayesianSlepianEncoder	class	Slepian encoder with NumPyro sites for cap radius and centre.

Spectral-Normalized Gaussian Process (SNGP) output layer. (pyrox_nn._sngp)

Name	Kind	What it does
RandomFeatureGaussianProcess	class	SNGP output layer (Liu et al., 2020).

Deep Variational SSGP — Bayesian wrapper around the geonnax core. (pyrox_nn._vssgp)

Name	Kind	What it does
DeepVSSGP	class	Deep Random Feature Expansion for Variational SSGP (Cutajar et al. 2017).

Re-exported from upstream

Name	Kind	What it does
AbstractConditioner	class	Re-exported from geonnax.
AffineModulation	class	Re-exported from geonnax.
BasisTransform	class	Re-exported from geonnax.
Cartesian3DEncoder	class	Re-exported from geonnax.
ConcatConditioner	class	Re-exported from geonnax.
ConditionedINR	class	Re-exported from geonnax.
CyclicEncoder	class	Re-exported from geonnax.
Deg2Rad	class	Re-exported from geonnax.
FourierFilter	class	Re-exported from geonnax.
FourierNet	class	Re-exported from geonnax.
GaborFilter	class	Re-exported from geonnax.
GaborNet	class	Re-exported from geonnax.
GeoFNO	class	Re-exported from geonnax.
HybridSphericalSlepianEncoder	class	Re-exported from geonnax.
HyperLinear	class	Re-exported from geonnax.
HyperSIREN	function	Re-exported from geonnax.
LaplaceRandomFeatureCovariance	class	Re-exported from geonnax.
LearnedDeformation	class	Re-exported from geonnax.
LonLatScale	class	Re-exported from geonnax.
ManifoldNO	class	Re-exported from geonnax.
NCPContinuousPerturb	class	Re-exported from geonnax.
OrthogonalRandomFeatures	class	Re-exported from geonnax.
SIREN	class	Re-exported from geonnax.
SirenDense	class	Re-exported from geonnax.
SlepianEncoder	class	Re-exported from geonnax.
SphericalHarmonicEncoder	class	Re-exported from geonnax.
cyclic_encode	function	Re-exported from geonnax.
deg2rad	function	Re-exported from geonnax.
fourier_features	function	Re-exported from geonnax.
interaction_features	function	Re-exported from geonnax.
lonlat_scale	function	Re-exported from geonnax.
lonlat_to_cartesian3d	function	Re-exported from geonnax.
mfn_forward	function	Re-exported from geonnax.
seasonal_features	function	Re-exported from geonnax.
seasonal_frequencies	function	Re-exported from geonnax.
spherical_harmonic_encode	function	Re-exported from geonnax.
standardize	function	Re-exported from geonnax.
unstandardize	function	Re-exported from geonnax.

pyrox_nn.api

BNFEstimator family — sklearn-style facade over the BNF stack. (pyrox_nn.api._bnf)

Name	Kind	What it does
BNFEstimator	class	MAP-fit Bayesian Neural Field estimator.
BNFEstimatorMLE	class	BNF estimator with prior_weight=0 (MLE).
BNFEstimatorVI	class	BNF estimator using mean-field VI via ensemble_vi.
FittedBNF	class	Output of BNFEstimator.fit.

Minimal sklearn-style immutable Estimator facade. (pyrox_nn.api._estimator)

Name	Kind	What it does
EstimatorBase	class	Immutable configuration for a probabilistic regression model.
FittedEstimator	class	Output of EstimatorBase.fit.

NeuralFieldEstimator — fit any Bayesian neural field with ensembles. (pyrox_nn.api._field)

Name	Kind	What it does
FittedNeuralField	class	Output of NeuralFieldEstimator.fit.
NeuralFieldEstimator	class	Fit any Bayesian neural field with ensemble MAP or VI.

pyrox_nn.preprocessing

Name	Kind	What it does
SpatiotemporalFit	class	Immutable bundle of fitted feature layers + time-encoding scalars.
encode_time_column	function	Convert a pandas time column into a unit-scale JAX float array.
fit_spatiotemporal	function	Build a complete SpatiotemporalFit from a DataFrame.
fit_standardization	function	Build a Standardization layer from per-column mean / std.

pyrox-lgm

pyrox_lgm

Areal (graph) components: Besag (ICAR), BYM2, proper CAR and Leroux (P7). (pyrox_lgm._components._areal)

Name	Kind	What it does
BYM2	class	BYM2 (Riebler et al., 2016), a mixture of unstructured and ICAR effects.
Besag	class	Intrinsic CAR (ICAR / Besag): precision \(\tau R\), \(R = D - W\).
CAR	class	Proper CAR: precision \(\tau(D - \rho W)\), \(\rho \in [0, 1)\).
Leroux	class	Leroux CAR: precision \(\tau(\rho R + (1 - \rho) I)\), \(\rho \in [0, 1]\).

Base class for latent components (P7). (pyrox_lgm._components._base)

Name	Kind	What it does
AbstractComponent	class	A latent field with hyperpriors and a projector to observations.

Combinators: separable Kronecker products and replicates of components (P7). (pyrox_lgm._components._combinators)

Name	Kind	What it does
Kronecker	class	Separable product \(Q = Q_\text{main} \otimes Q_\text{group}\) (R-INLA group).
Replicate	class	n_rep independent copies of a component, sharing its hyperparameters.

Generic components from a user-supplied structure (P7). (pyrox_lgm._components._generic)

Name	Kind	What it does
Generic	class	A field with precision \(\tau R\) for a user-supplied structure \(R\).

SPDE (Matérn) components on a finite-element mesh or a regular grid (P7). (pyrox_lgm._components._spde)

Name	Kind	What it does
SPDE	class	Matérn field via the SPDE, on a triangle mesh or a regular grid.

Temporal (and unstructured) components: IID, RW1, RW2 and AR1 (P7). (pyrox_lgm._components._temporal)

Name	Kind	What it does
AR1	class	Stationary AR(1), \(x_t = \rho x_{t-1} + \varepsilon_t\).
IID	class	Independent effects \(x_i \sim \mathcal N(0, 1/\tau)\), \(i < n\).
RW1	class	First-order random walk, \(x_{t+1} - x_t \sim \mathcal N(0, 1/\tau)\).
RW2	class	Second-order random walk, \(x_{t+1} - 2x_t + x_{t-1} \sim \mathcal N(0, 1/\tau)\).

Model diagnostics without refits: DIC, WAIC, CPO and PIT (P9). (pyrox_lgm._diagnostics)

Name	Kind	What it does
Diagnostics	class	Leave-one-out and information criteria of an INLAResult.
diagnostics	function	DIC, WAIC, CPO and PIT of an INLAResult, with no refits.

f(...): R-INLA-style construction of a component bound to a data column (P9). (pyrox_lgm._formula)

Name	Kind	What it does
f	function	Build a latent component named after the data column it indexes.

The inla() driver (P8). (pyrox_lgm._inla)

Name	Kind	What it does
inla	function	Integrated nested Laplace approximation of a latent Gaussian model.

Observation models of an LGM (P8). (pyrox_lgm._likelihood)

Name	Kind	What it does
AbstractObservation	class	y \| eta for an LGM, with its own hyperparameters.
Bernoulli	class	y ~ Bernoulli(sigmoid(eta)), y in {0, 1}.
Binomial	class	y ~ Binomial(n, sigmoid(eta)) with n = data["n_trials"].
Gaussian	class	y ~ N(eta, 1/prec); prec under PCPrecision(1, 0.01).
NegativeBinomial	class	y ~ NegBin(mean exp(eta), size); size under Gamma(1, 0.1).
Poisson	class	y ~ Poisson(exp(eta)); put the log exposure in data["offset"].

The latent Gaussian model, LGM (P8). (pyrox_lgm._model)

Name	Kind	What it does
FixedEffects	class	Fixed effects \(\beta \sim \mathcal N(0, \lambda^{-1} I)\), columns of \(X\).
LGM	class	A latent Gaussian model: components, fixed effects and an observation model.

Penalised-complexity (PC) priors (Simpson et al., 2017). (pyrox_lgm._priors._pc)

Name	Kind	What it does
PCAR1Rho	class	PC prior on the lag-one correlation \(\rho\) of an AR(1) process.
PCBYM2Phi	class	PC prior on the BYM2 mixing parameter \(\phi \in (0, 1)\).
PCMatern	class	Joint PC prior on the range \(\rho\) and marginal sd \(\sigma\) of a Matérn field.
PCPrecision	class	PC prior on the precision \(\tau\) of a Gaussian effect.
StructureSpectrum	class	Spectral measure of \(R_\ast^{+}\), the generalised inverse of a structure.
structure_spectrum	function	The spectrum of \(R^{+}\) that PCBYM2Phi needs, computed once.

INLAResult: posterior marginals and the marginal likelihood (P8). (pyrox_lgm._result)

Name	Kind	What it does
INLAResult	class	What inla() returns.
Summary	class	Posterior summary: mean, sd and the 2.5 / 50 / 97.5 % quantiles.

Listed at gaussx 0.6.5, kernellib 0.0.17, geonnax 0.0.11.

Upstream: gaussx

Structured linear algebra and Gaussians: solve / logdet / Cholesky on structured operators, Gaussian densities and KLs, GP conditioning, Kalman filters, quadrature. Use it instead of jnp.linalg or cho_solve on a covariance.

gaussx

Name	Kind	What it does
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.
AbstractMultivariateNormal	class	Common surface of MultivariateNormal and MultivariateNormalPrecision.
AbstractPreconditioner	class	Protocol for preconditioners producing an approximate inverse.
AbstractSketch	class	A random sketching matrix \(S \in \mathbb{R}^{d \times m}\), sampled once.
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	CG solve and SLQ logdet with BBMM's defaults (Gardner et al. 2018).
BYM2GMRF	class	The BYM2 pair \((b, u^\ast)\) (Riebler et al., 2016), with its exact density.
BernoulliLikelihood	class	Bernoulli likelihood with logit link.
BinomialLikelihood	class	Binomial likelihood with logit link.
BlockDiag	class	Block diagonal operator diag(A₁, A₂, …, Aₖ).
BlockTriDiag	class	Block-tridiagonal operator with a symmetric off-diagonal band.
CGSolver	class	Iterative CG solver with stochastic log-determinant.
CUR	class	CUR decomposition \(A \approx C\,U\,R\), \(C = A_{:,J}\), \(R = A_{I,:}\).
CapacitanceSolver	class	Solve a base system subject to homogeneous point constraints.
Chandrupatla	class	Chandrupatla's hybrid inverse-quadratic / bisection root finder.
ColumnID	class	Column interpolative decomposition \(A \approx A_{:,J}\,X\).
ComposedSolver	class	Mix-and-match solve and logdet from different strategies.
ConstantSDE	class	State-space representation of a constant kernel.
ConstrainedGMRF	class	A GaussianMRF conditioned on hard linear constraints \(A_cx = e\).
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.
DiagonalizedOperator	class	Linear operator A x = V⁻¹(Λ ⊙ V x) given by a transform pair.
EigenDecomposition	class	Precomputed eigendecomposition A = V diag(λ) V⁻¹ of a square matrix.
EmissionModel	class	Observation (emission) model wrapping a linear observation matrix.
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).
GaussianMRF	class	Gaussian Markov random field \(x \sim \mathcal N(\mu, Q^{-1})\).
GaussianSites	class	Time-varying Gaussian likelihood sites in natural parameterization.
GaussianSketch	class	Gaussian sketch: i.i.d. entries \(S_{ij} \sim \mathcal{N}(0, 1/d)\).
GaussianState	class	Gaussian distribution as (mean, covariance operator) pair.
HeteroscedasticGaussianLikelihood	class	Heteroscedastic Gaussian likelihood with input-dependent noise.
IndefiniteSLQLogdet	class	Stochastic log\|det(A)\| for symmetric (possibly indefinite) operators.
IntegratedWienerSDE	class	State-space representation of an integrated Wiener process.
InterpolatedOperator	class	Structured Kernel Interpolation: K \approx W K_{uu} W^T.
IntrinsicGMRF	class	Improper \(\mathcal N(\mu, (\tau R)^+)\) on the complement of \(\ker R\).
JacobiPreconditioner	class	Diagonal preconditioner M^{-1} = diag(1 / \|diag(A)\|).
KeyedSolver	class	Bind a PRNG key to a strategy's stochastic log-determinant (gh-384).
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.
LGSSM	class	Linear-Gaussian state-space model as a density over (T, M).
LGSSMFactory	class	Callable obs_mask -> MaskedLGSSM, for conditional use.
LOOResult	class	Result of leave-one-out cross-validation.
LOVECache	class	Cached rank-k Lanczos approximation of K^{-1}.
LSMRSolver	class	LSMR iterative least-squares solver (Fong & Saunders 2011).
LaplaceResult	class	The Gaussian approximation \(\mathcal N(\hat x, H^{-1})\) at the mode.
LinearizedSDE	class	Linear surrogate \(dx = (A_t x + b_t)\,dt + Q_t^{1/2}\,dW\) of an SDE.
LineaxSolver	class	A lineax solver as a gaussx solve strategy.
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.
MarkovGaussian	class	Chain-Markov Gaussian over \(x_{0:T-1}\) as a density over (T, d).
MaskedLGSSM	class	LGSSM whose log_prob marginalises unobserved channels exactly.
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.
MultivariateNormal	class	Multivariate normal parameterized by a lineax linear operator.
MultivariateNormalPrecision	class	Multivariate normal parameterized by a precision (inverse covariance) operator.
NegativeBinomialLikelihood	class	Negative-binomial likelihood with log link (NB2, a gamma-Poisson mixture).
NystromLogdet	class	Nyström-preconditioned log-determinant of \(A + \mu I\) (Wenger et al., 2022).
NystromPreconditioner	class	Randomized Nyström preconditioner for A + μ I.
OperatorPreconditioner	class	Use an externally supplied approximate inverse as a preconditioner.
OrthonormalSketch	class	Gaussian sketch with orthonormal rows, \(S S^\top = I_d\).
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 (and factor) for amortized predictions.
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.
RowSamplingSketch	class	Weighted row-sampling sketch.
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.
SRHTSketch	class	Subsampled randomized Hadamard transform sketch.
SVDLowRankUpdate	class	Deprecated subclass of LowRankUpdate with orthonormal=True.
SketchAndPrecondLSMR	class	Sketch-and-precondition LSMR for tall least squares (Rokhlin & Tygert, 2008).
SoftmaxLikelihood	class	Softmax (categorical) likelihood for multi-class classification.
SparseCholeskyFactor	class	Sparse Cholesky factor P Q Pᵀ = L Lᵀ on a static symbolic pattern.
SparseCholeskySolver	class	Exact solve / logdet / diag_inv through a sparse Cholesky factor.
SparseOperator	class	Sparse matrix with a static SparsityPattern and traced values.
SparseSignSketch	class	Sparse sign sketch (sparse Johnson-Lindenstrauss transform).
SparsityPattern	class	Static, hashable sparsity pattern of a (m, n) matrix.
SpectralFunction	class	f(B) for a symmetric B with a factored eigenbasis.
StudentTLikelihood	class	Student-t likelihood for robust regression.
SumKronecker	class	Deprecated alias for SumOfKroneckers.
SumKroneckerSqrt	class	Deprecated alias for SumOfKroneckersSqrt (gh-297).
SumOfKroneckers	class	Sum of Kronecker products Σ_k A_k \otimes B_k.
SumOfKroneckersSqrt	class	Lazy Lanczos square-root operator for SumOfKroneckers covariances.
SumSDE	class	Sum of SDE kernels via block-diagonal composition.
SymbolicCholesky	class	Symbolic Cholesky factor of a symmetric sparsity pattern.
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.
ar1_precision	function	Precision of a stationary AR(1) process with marginal precision τ.
as_diagonalized	function	Express operator as one DiagonalizedOperator, if possible.
as_linear_operator	function	Wrap a raw matvec callable as a tagged lineax operator.
base_conditional	function	Deprecated: use sparse_conditional (note K_xz = K_mn.T).
besag_structure	function	Validate a graph Laplacian as a Besag (ICAR) structure matrix.
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 K_y alpha = y and cache the result.
bym2_precision	function	Joint precision of the BYM2 pair (b, u*) (Riebler et al., 2016).
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.
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	function	(Block-)circulant operator from its first column.
circulant_from_symbol	function	(Block-)circulant operator from its DFT symbol.
collapsed_elbo	function	Collapsed ELBO (Titsias bound) for sparse GP regression.
column_id	function	Randomized column interpolative decomposition (Voronin & Martinsson, 2017).
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	Two-filter fusion at t from a filtered state and a backward message.
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.
cur	function	Randomized CUR decomposition (Voronin & Martinsson, 2017, §4).
cvi_update_sites	function	Natural gradient update for CVI sites.
damped_natural_update	function	Damped update in natural parameter space.
dare	function	Steady-state Kalman filter covariance and gain (the filtering DARE).
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.
discretize_mfd	function	Discretise a linear SDE by matrix-fraction decomposition.
discretize_mfd_sequence	function	Vectorised gaussx.discretize_mfd over a vector of time steps.
eig	function	Compute eigenvalues and eigenvectors.
eigh_generalized	function	Solve \(A v = \lambda B v\) for symmetric \(A\) and symmetric PSD \(B\).
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	Moments of the scalar EP tilted distribution, elementwise over a batch.
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_mean_chol	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.
fem_matrices	function	Lumped mass C̃ and stiffness G of P1 elements on a triangle mesh.
fem_projector	function	Observation matrix A of P1 interpolation: (A w)_k = Σ_i ψ_i(s_k) w_i.
fifth_order_cubature_points	function	Fifth-order fully symmetric cubature points and weights.
fisher_info	function	Fisher information with respect to the mean, \(F_\mu = \Sigma^{-1}\).
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(first ‖ second): KL[q(u) \|\| p(u)] with p(u) = N(0, K).
gauss_newton_precision	function	Gauss-Newton precision matrix J^T J, optionally plus a prior.
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_kl	function	KL(first ‖ second): KL(p \|\| q) between two multivariate normals.
gaussian_log_prob	function	Multivariate normal log-probability.
generalized_variance_scale	function	Generalized variance of an intrinsic GMRF (Sørbye & Rue, 2014).
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.
hadamard_transform	function	Unnormalized Walsh-Hadamard transform along the last axis.
haversine_distance	function	Pairwise great-circle (haversine) distances on a sphere.
hutchinson_hessian_diag	function	Stochastic Hessian diagonal via Hutchinson with Rademacher probes.
iid_precision	function	Precision τ I of n independent effects with variance 1/τ.
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_eigen_reducible	function	Whether a sum of Kronecker products has the exact two-term reduction.
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(first ‖ second): KL(q \|\| p) via the Bregman form on natural parameters.
kl_standard_normal	function	KL(first ‖ second): 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
laplace_mode	function	Mode, Hessian and log-marginal of the Laplace approximation, in precision form.
leave_one_out_cv	function	LOO-CV via the bordered-system identity.
linear_solve	function	Solve A x = b through the unified front door.
linearize_sde	function	Optimal linear drift of a nonlinear SDE along a Gaussian path.
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 a Lanczos approximation of K^{-1} for fast variance.
love_residual	function	Relative solve residual of a LOVE cache for one cross-covariance.
love_variance	function	Approximate k_*^\top K^{-1} k_* from a LOVE cache.
love_variance_error_bound	function	Upper bound on the variance over-estimate of 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.
matern_spde_params	function	SPDE parameters (κ, τ, α) of a Matérn field.
matheron_update	function	Posterior samples via Matheron's-rule correction.
mean_chol_to_expectation	function	Convert mean/variance (Cholesky) to expectation parameters.
mean_chol_to_natural	function	Convert mean/variance (Cholesky) to natural parameters.
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.
mixture_quantile	function	Invert a CDF at one or more quantile levels by bracketed root finding.
mixture_quantile_gaussian_approx	function	Quantiles of the Gaussian matched to an ensemble's first two moments.
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_chol	function	Convert natural parameters to mean/variance (Cholesky).
natural_to_mean_cov	function	Convert natural parameters to mean/covariance (operator form).
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.
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.
product_operator	function	Lazy matmul (A B) v = A (B v).
project	function	Compute A_X = K_XZ @ K_ZZ^{-1} via Cholesky solve.
pseudo_logdet	function	Log pseudo-determinant: the log of the product of the non-zero eigenvalues.
qb	function	Randomized QB factorisation \(A \approx QB\) with \(B = Q^\top A\).
quadratic_form	function	Compute x^T A^{-1} x via a single solve.
randomized_eigh	function	Partial eigendecomposition of a symmetric operator by Rayleigh-Ritz.
randomized_nystrom	function	Randomized Nyström approximation \(\hat A = U\hat\Lambda U^\top\) of a PSD \(A\).
randomized_svd	function	Truncated SVD \(A \approx U \operatorname{diag}(s) V^\top\) by randomized QB.
range_finder	function	Orthonormal basis \(Q\) for the dominant range of \(A\).
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⁻¹.
rp_cholesky	function	Randomly pivoted partial Cholesky of a PSD matrix A.
rts_interpolate	function	Smoothed marginal at an unobserved t from filtered and smoothed states.
rts_smoother	function	Rauch-Tung-Striebel backward smoother.
rw1_structure	function	Structure matrix R = D₁ᵀ W D₁ of a first-order random walk.
rw2_structure	function	Structure matrix R = D₂ᵀ D₂ of a second-order random walk.
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.
scaled_operator	function	Lazy scalar multiply (c A) v = c (A v).
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.
sde_kl_divergence	function	Drift path-KL between a linear and a nonlinear SDE along a Gaussian path.
selected_inverse	function	The band of Σ = Q⁻¹ for a symmetric positive-definite BlockTriDiag.
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.
sketch_and_solve	function	Approximate least squares from a sketch.
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.
sparse_cholesky	function	Sparse Cholesky factorisation of a symmetric positive-definite operator.
sparse_conditional	function	Sparse-GP conditional q(f_x) = \int p(f_x \| u) q(u) du.
spde_precision	function	SPDE precision Q_α = τ² K (C̃⁻¹K)^{α−1} with K = κ²C̃ + G.
spde_precision_grid	function	SPDE precision on a regular grid: a function of a Kronecker sum.
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_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].
sum_of_kroneckers_sample	function	Sample from 𝒩(0, op) with matrix-free Lanczos square roots.
sum_operator	function	Lazy sum (A + B + …) v = A v + B v + ….
sumkronecker_sample	function	Deprecated alias for sum_of_kroneckers_sample (gh-297).
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.
symbolic_cholesky	function	Symbolic analysis of a sparse Cholesky factorisation (host, cached).
symmetric_tag	constant	Re-exported from lineax.
symmetrize	function	Symmetrize the trailing two axes: 0.5 * (X + X^T).
theta_design	function	Integration points and log-weights over hyperparameters \(\theta\).
tikhonov_augment	function	Observation augmentation for Tikhonov-regularised EKI (TEKI).
to_expectation	function	Deprecated alias of to_mean_cov.
to_mean_cov	function	Convert natural parameters to the mean and covariance.
to_natural	function	Deprecated alias of gaussx.mean_cov_to_natural.
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_solve	function	Solve a tridiagonal system A x = d.
tridiagonal_solve_batched	function	Solve independent tridiagonal systems over leading batch dimensions.
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	Uncollapsed (SVGP) ELBO for a Gaussian likelihood.
variational_elbo_mc	function	Monte Carlo ELBO for non-conjugate likelihoods.
vb_mean_correction	function	Low-rank VB correction of the Laplace mean (Van Niekerk & Rue, 2024).
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.

Upstream: kernellib

Kernels (AbstractKernel, stationary and dot-product families), kernel operators and the pure kernel functions in kernellib.functional. pyrox-gp's kernels wrap these.

kernellib

Name	Kind	What it does
AbstractEstimator	class	Configure, fit, predict.
AbstractFeatureMap	class	A finite-dimensional approximation k(x, x') ≈ φ(x)ᵀ φ(x').
AbstractGraph	class	Weighted, undirected graph without self-loops.
AbstractKernel	class	Anything that produces a Gram matrix.
AbstractPointwiseKernel	class	A kernel defined by k(x, x').
AbstractStationaryKernel	class	k(x, x') = variance * shape(\|\|(x - x') / lengthscale\|\|²).
ActiveDims	class	Apply a kernel to a subset of input dimensions.
CKAAccumulator	class	Running mini-batch CKA (Nguyen, Raghu & Kornblith, 2021).
Constant	class	Constant kernel, k(x, x') = sigma^2.
Cosine	class	Cosine kernel, sigma^2 cos(2 pi \|\|x - x'\|\| / p).
Derivative	class	\(\partial_{x_i} \partial_{x'_j} k(x, x')\) for fixed input dimensions.
DerivativeIndexed	class	Joint covariance of a GP's values and partial derivatives.
Distance	class	Distance-induced kernel, k(x, x') = σ² (\|x\|^a + \|x'\|^a - \|x - x'\|^a) / 2.
EigenPro	class	Kernel regression by EigenPro-preconditioned SGD (Ma & Belkin, 2017).
EigenProPreconditioner	class	Spectral correction state for EigenPro kernel SGD.
Falkon	class	Nyström kernel ridge regression solved by Falkon (Rudi et al., 2017).
FalkonInfo	class	Convergence report of a falkon_solve run.
FalkonPreconditioner	class	The two upper-triangular Cholesky factors of the Falkon preconditioner.
FastFoodFeatures	class	FastFood random features (Le, Sarlós & Smola, 2013) for any kernel with a spectral sampler.
FastFoodParams	class	The random structure of a FastFood feature map.
FeatureKernel	class	Inner product of features, \(k(x, x') = \langle \phi(x), \phi(x') \rangle\).
Graph	class	Sparse weighted graph: a static topology and traced edge weights.
GraphTopology	class	Static, hashable edge list of an undirected graph without self-loops.
GridGraph	class	Regular lattice graph that never stores its edges.
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.
KNNGraph	class	A directed k-nearest-neighbour graph, self excluded.
KRR	class	Kernel ridge regression.
KernelLocalityPreservingProjections	class	Kernel locality preserving projections.
KernelOperator	class	Kernel matrix operator with efficient first-order autodiff.
KernelPCA	class	Kernel principal component analysis.
KernelSchrodingerProjections	class	Kernel Schrödinger eigenmap projections.
LaplaceEigenfunctionFeatures	class	Hilbert-space (HSGP) approximation of a stationary kernel on a box.
LaplacianEigenmaps	class	Laplacian eigenmaps (Belkin & Niyogi, 2003) of the k-NN graph of X.
Linear	class	Linear kernel, k(x, x') = sigma^2 x^T x' + b.
LocalityPreservingProjections	class	Locality preserving projections (He & Niyogi, 2003).
Matern	class	Matern kernel with closed-form nu in {0.5, 1.5, 2.5}.
Modulated	class	A kernel with an input-dependent amplitude, \(k(x, x') = a(x)\, k_0(x, x')\, a(x')\).
NystromFeatures	class	Nyström features (Williams & Seeger, 2001) for any kernel.
OrthogonalRandomFeatures	class	Orthogonal random features (Yu et al., 2016).
Periodic	class	Periodic (MacKay) kernel on the Euclidean distance.
Periodised	class	Make any kernel periodic: apply it to inputs wrapped onto circles.
PermutationTestResult	class	Outcome of permutation_test.
Polynomial	class	Polynomial kernel, k(x, x') = sigma^2 (x^T x' + b)^d.
Product	class	Elementwise product of kernels, k(x, x') = prod_i k_i(x, x').
RBF	class	Radial basis function (squared exponential) kernel.
RandomFourierFeatures	class	Random Fourier features (Rahimi & Recht, 2007).
RationalQuadratic	class	Rational quadratic kernel, a scale mixture of RBFs.
Residual	class	What an approximation misses, \(k(x, x') - \tilde k(x, x')\).
Scaled	class	A kernel times a scalar, k(x, x') = scale * k_0(x, x').
SchrodingerEigenmapProjections	class	Schrödinger eigenmap projections (SEP): linear Schrödinger eigenmaps.
SchrodingerEigenmaps	class	Schrödinger eigenmaps (Czaja & Ehler, 2013) of the k-NN graph of X.
Shift	class	Input warp x -> x - offset; offset is scalar or per dimension.
Stretch	class	Input warp x -> x / scale; scale is scalar or per dimension.
Sum	class	Sum of kernels, k(x, x') = sum_i k_i(x, x').
TaylorStatistics	class	The sides and angle of a kernel Taylor-diagram triangle.
Warped	class	Apply a kernel to warped inputs, k(x, x') = k_0(w(x), w(x')).
White	class	White-noise kernel, sigma^2 delta(x, x').
adjacency_matrix	function	Symmetric weighted adjacency matrix of a k-NN graph.
barrier_potential	function	Diagonal potential that is strength at indices and 0 elsewhere.
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.
cka	function	Centred kernel alignment, \(\mathrm{HSIC}(x, y) / \sqrt{\mathrm{HSIC}(x, x)\,\mathrm{HSIC}(y, y)}\).
combine_potentials	function	Several Schrödinger potentials as one: \(\sum_k \alpha_k c_k V_k\).
commute_time_kernel	function	Commute-time kernel \(K = L^{+}\), the Laplacian's pseudo-inverse.
cosine_graph_kernel	function	Inverse-cosine kernel \(K = \cos(\pi L / 4)\) (Smola & Kondor, 2003).
delaunay_graph	function	The edge graph of the Delaunay triangulation of X.
derivative_inputs	function	Rows [x, i] for DerivativeIndexed: values, then partials.
diffusion_kernel	function	Diffusion (heat) kernel \(K = \exp(-\beta L)\) (Kondor & Lafferty, 2002).
distance_correlation_squared	function	Squared distance correlation (Székely et al., 2007), cka.
distance_covariance_squared	function	Squared distance covariance (Székely et al., 2007), 4 * hsic.
draw_rff_cosine_basis	function	Draw (variance, lengthscale, omega, phase, weights) for a kernel.
edge_weights	function	Reweight a graph's edges by a kernel on the nodes' features.
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.
energy_distance	function	Energy distance (Székely & Rizzo, 2004), 2 * mmd_squared.
estimate_lengthscale	function	A data-driven lengthscale for an RBF-type kernel.
evaluate_rff_cosine_paths	function	Evaluate the zero-mean RFF prior path(s) at inputs X.
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.
fastfood_features	function	FastFood feature map, \(\Phi(x) = D^{-1/2}\,[\cos(Vx), \sin(Vx)]\).
fastfood_frequencies	function	The dense frequency matrix \(V\) that FastFood applies implicitly.
fastfood_operator	function	FastFood approximation of the RBF Gram matrix, \(K \approx \Phi\Phi^\top\).
fastfood_params	function	Draw the random structure for a FastFood RBF feature map.
gabriel_graph	function	The Gabriel graph of X.
gamma_to_lengthscale	function	lengthscale = 1 / sqrt(2 gamma), the inverse of lengthscale_to_gamma.
graph_from_adjacency	function	The sparse graph of a dense, symmetric adjacency matrix.
graph_from_edges	function	The sparse graph of an edge list: the seam for other graph libraries.
graph_from_neighbors	function	The sparse, symmetric graph of a k-nearest-neighbour relation.
graph_laplacian	function	Graph Laplacian of a symmetric adjacency matrix.
graph_null_space	function	Orthonormal basis of the Laplacian's null space.
grid_graph	function	The lattice graph of an array of shape shape.
hadamard_transform	function	Re-exported from gaussx.
hsic	function	Hilbert-Schmidt Independence Criterion between paired samples.
hsic_penalty	function	The HSIC penalty operator \(M = H K_s H / n^2\) on protected attributes.
implicit_cross_kernel	function	Create a matrix-free rectangular cross-kernel operator.
kernel_alignment	function	Uncentred kernel alignment, $\langle K, L\rangle_F / (\
knn_graph	function	The k-nearest-neighbour graph of X.
label_potential	function	Non-diagonal potential pulling together points that share a label.
laplacian_eigenmap	function	Laplacian eigenmap of a weighted graph: smallest solutions of \(L y = \lambda D y\).
laplacian_eigpairs	function	The n smallest eigenpairs of a graph Laplacian.
laplacian_penalty	function	The graph-smoothness penalty operator \(M = L / n^2\).
lengthscale_grid	function	Log-spaced lengthscales within decades orders of magnitude of center.
lengthscale_to_gamma	function	gamma = 1 / (2 lengthscale^2), the scikit-learn RBF parameter.
matern_graph_kernel	function	Graph Matérn kernel \(K = U \Phi(\Lambda) U^\top\) (Borovitskiy et al., 2021).
mesh_graph	function	The edge graph of a triangle mesh.
mmd_squared	function	Squared maximum mean discrepancy between samples X and Y.
n_components_graph	function	Number of connected components (an isolated node is one).
nearest_neighbors	function	The n_neighbors nearest other points of every point in X.
nystrom_kernel	function	The Nyström kernel \(k_Z(x, x') = k(x, Z) K_{ZZ}^{-1} k(Z, x')\).
nystrom_operator	function	Nystrom low-rank kernel approximation.
permutation_test	function	Permutation p-value for a statistic that is large under the alternative.
radius_graph	function	The graph joining points within radius of each other.
radius_neighbors	function	The other points within radius of every point, at most max_neighbors of them.
random_walk_kernel	function	\(p\)-step random-walk kernel \(K = (aI - L)^p\) (Smola & Kondor, 2003).
regularized_laplacian_kernel	function	Regularised Laplacian kernel \(K = (I + \sigma^2 L)^{-1}\).
relative_neighborhood_graph	function	The relative neighbourhood graph of X.
rff_operator	function	Random Fourier Features kernel approximation.
schrodinger_eigenmap	function	Schrödinger eigenmap: smallest solutions of \((L + \alpha V) y = \lambda D y\).
select_landmarks	function	Indices into X of n_landmarks distinct landmarks.
spatial_spectral_graph	function	Spatial-spectral graph for images (Cahill, Czaja & Messinger, 2014).
spatial_spectral_potential	function	Spatial-spectral potential for images (Cahill, Czaja & Messinger, 2014).
structure_matrix	function	The Besag (ICAR) structure matrix \(R = L\), optionally BYM2-scaled.
taylor_statistics	function	Radii, correlation and distance for a kernel Taylor diagram.
to_cross_operator	function	K(X1, X2) as a rectangular linear operator.
to_operator	function	K(X, X) (+ noise * I) as a symmetric, PSD-tagged linear operator.

kernellib.functional

Name	Kind	What it does
center_cross_kernel	function	Centre a cross-kernel matrix with the training Gram's statistics.
center_kernel	function	Center a kernel matrix: H K H.
centering_operator	function	Centering matrix H = I - (1/n) \mathbf{1}\mathbf{1}^T.
cka	function	Centered kernel alignment.
constant_kernel	function	Constant kernel.
cosine_kernel	function	Cosine kernel.
distance_kernel	function	Distance-induced kernel (Sejdinovic et al., 2013).
graph_heat_spectrum	function	Heat-kernel spectrum \(\Phi(\lambda) = \sigma^2 \exp(-\ell^2 \lambda / 2)\).
graph_matern_spectrum	function	Graph Matérn spectrum \(\Phi(\lambda) = \sigma^2 (2\nu/\ell^2 + \lambda)^{-\nu}\).
hsic	function	Hilbert-Schmidt Independence Criterion from two kernel matrices.
kernel_add	function	Pointwise sum of two already-evaluated Gram matrices.
kernel_mul	function	Pointwise (Hadamard) product of two already-evaluated Gram matrices.
linear_kernel	function	Linear kernel.
matern_kernel	function	Matern kernel with closed-form nu in {1/2, 3/2, 5/2}.
mmd_squared	function	Biased squared Maximum Mean Discrepancy.
periodic_kernel	function	Periodic (MacKay) kernel.
polynomial_kernel	function	Polynomial kernel.
rational_quadratic_kernel	function	Rational quadratic kernel.
rbf_kernel	function	Radial basis function (squared exponential) kernel.
stable_rbf_kernel	function	RBF (squared exponential) kernel with mixed-precision stability.
white_kernel	function	White-noise kernel.

Upstream: geonnax

Deterministic Equinox network cores, encoders and basis functions. pyrox-nn wraps a geonnax core and swaps its parameters for sites.

geonnax

Name	Kind	What it does
AbstractConditioner	class	Duck-typed protocol for (h, z) -> y conditioning layers.
AffineModulation	class	Feature-wise Linear Modulation (FiLM): y = γ(z) ⊙ h + β(z).
AntiAliasedActivation	class	A pointwise nonlinearity between band-limited function spaces.
Attention	class	Pre-norm multi-head self-attention with cosine (QK-normalised) similarity.
BasisSpectralConv	class	Spectral convolution in a precomputed basis: per-mode channel mixing.
BasisTransform	class	Frozen analysis / synthesis pair for a precomputed basis.
Block	class	Conv → GroupNorm → SiLU, the residual-block primitive.
CNO	class	Convolutional Neural Operator: lift → U-Net of band-limited blocks → project.
CNOBlock	class	Convolution then band-limited activation: the CNO I, D and U blocks.
CNOResBlock	class	The CNO residual (R) block, \(v \mapsto v + K_2\,\Sigma\,K_1 v\).
CPTensor	class	Canonical-Polyadic (CP / PARAFAC) factorization.
Cartesian3DEncoder	class	Lift a single lon/lat coordinate onto the unit sphere \(S^2\).
ConcatConditioner	class	Concatenate h and z then apply a single Linear.
ConditionedINR	class	Wrap an inner network's per-layer activations with conditioners.
ConvNeXtBlock	class	ConvNeXt-V2 block: depthwise conv → norm → MLP with GRN (Woo et al., 2023).
ConvolutionalNeuralOperator	class	Convolutional Neural Operator: lift → U-Net of band-limited blocks → project.
CyclicEncoder	class	Encode a single periodic input as concatenated cos/sin features.
DeepVSSGPCore	class	Deterministic deep RFF spectral core.
Deg2Rad	class	Element-wise degrees-to-radians conversion.
DenseRank1	class	Rank-1 ensemble dense layer (deterministic BatchEnsemble core).
DenseTensor	class	A full, un-factorised tensor.
DomainPadding	class	Zero-pad spatial axes by a fraction of their size, then crop back.
Downsample	class	Memory-efficient factor-2 downsample: space-to-depth, then 1x1 conv.
FNO	class	Fourier Neural Operator: lifting → Fourier blocks → projection.
FNOBlock	class	One Fourier layer: spectral conv + pointwise skip, then an activation.
FiLM	class	Feature-wise Linear Modulation (FiLM): y = γ(z) ⊙ h + β(z).
FourierFilter	class	Single Fourier filter: \(g(x) = \sin(\Omega x + \varphi)\).
FourierNet	class	Multiplicative Fourier Filter Network (Fathony et al., ICLR 2021).
FourierNeuralOperator	class	Fourier Neural Operator: lifting → Fourier blocks → projection.
GaborFilter	class	Single Gabor filter: $g(x) = \sin(\Omega x + \varphi) \odot \exp(-\tfrac{\gamma}{2}\
GaborNet	class	Multiplicative Gabor Filter Network (Fathony et al., ICLR 2021).
GaussianFourierEncoder	class	Gaussian random Fourier-feature positional encoding (Tancik et al., 2020).
GaussianRegressionHead	class	Heteroscedastic Gaussian regression head (mean plus log-scale).
GeneratedSiren	class	Wrapper around a SIREN whose layers consume generated weights.
GeoContextEncoder	class	Encode geoscience metadata into a single context vector.
GeoFNO	class	Geo-FNO: a Fourier neural operator on per-sample point clouds.
GlobalResponseNorm	class	Global Response Normalization (Woo et al., 2023, "ConvNeXt-V2").
HeteroscedasticHead	class	Shared state and init for the FA-noise heteroscedastic layers.
HybridSphericalSlepianEncoder	class	Concatenate low-bandwidth global SHs with local Slepian cap modes.
HyperLinear	class	Generate a target Linear's (W, b) from z, then apply.
HyperSIREN	function	NIF-style ShapeNet/ParameterNet composite (Pan, Brunton, Kutz — JMLR 2023).
LNO	class	Local Neural Operator: pad → local core → crop, marched with jax.lax.scan.
LaplaceRandomFeatureCovariance	class	Laplace-approximation precision for an SNGP output head.
LayerNormEnsemble	class	Per-ensemble-member LayerNorm (deterministic core).
LearnedDeformation	class	A learned chart from the physical domain onto the latent torus (Geo-FNO).
LinearAttention	class	Pre-norm linear attention — O(n) in the number of spatial positions.
LinearChainCRF	class	Deterministic linear-chain CRF — forward algorithm plus Viterbi decoding.
LocalNeuralOperator	class	Local Neural Operator: pad → local core → crop, marched with jax.lax.scan.
LonLatScale	class	Affine-rescale a single lon/lat pair into [-1, 1].
MCSigmoidDenseFA	class	Heteroscedastic multi-label output layer (FA noise + sigmoid).
MCSoftmaxDenseFA	class	Heteroscedastic multi-class output layer (FA noise + softmax).
MSWT	class	Multi-Scale Wavelet Transformer operator (Wang et al., 2026).
MSWTBlock	class	Pre-norm wavelet-attention block: an attention residual, then an FFN residual.
ManifoldNO	class	Manifold Neural Operator: lifting → basis blocks → projection.
ManifoldNOBlock	class	One manifold-operator layer: basis spectral conv + pointwise skip.
MixtureOfGaussiansDenseHead	class	Deterministic mixture-density-network output head.
MultiHeadAttentionBE	class	Multi-head attention with BatchEnsemble rank-1 projections.
MultiScaleSIREN	class	Parallel frequency-banded SIRENs summed at the readout.
MultiScaleWaveletTransformer	class	Multi-Scale Wavelet Transformer operator (Wang et al., 2026).
NCPContinuousPerturb	class	Input perturbation for the Noise Contrastive Prior pattern.
NestedResidualUNet	class	A small residual U-Net used as one stage block (Qin et al., 2020, "U²-Net").
NeuralFieldMLP	class	The BayesNF spatiotemporal network as a deterministic equinox.Module.
OrthogonalRandomFeatures	class	Orthogonal Random Features (Yu et al., 2016) — variance-reduced RFF.
RandomFeatureGaussianProcess	class	SNGP output layer (Liu et al., 2020) — deterministic core.
Rank1ProjInit	class	Per-projection BatchEnsemble inits used by MultiHeadAttentionBE.
ResnetBlock	class	Two Blocks with squeeze-excitation and a residual skip.
SFNO	class	Spherical Fourier Neural Operator: lift → spherical blocks → project.
SIREN	class	Multi-layer sinusoidal representation network (Sitzmann et al., NeurIPS 2020).
SeparableINR	class	A rank-\(r\) separable neural field over \(A\) input axes.
SirenDense	class	Sine-activated dense layer: y = sin(ω · (W x + b)) or y = W x + b.
SirenLayerSpec	class	Static per-layer metadata shared by deterministic and Bayesian SIRENs.
SirenLayerType	function	Re-exported from typing.
SlepianEncoder	class	Deterministic Slepian-cap positional encoder on \(S^2\).
SpectralConv	class	Fourier spectral convolution over num_spatial_dims = len(n_modes) axes.
SpectralNormalization	class	Wrap an Equinox linear layer to be \(c\)-Lipschitz via power iteration.
SphericalFNOBlock	class	One spherical Fourier layer: spherical spectral conv + pointwise skip.
SphericalHarmonicEncoder	class	Real spherical-harmonic features on the unit 2-sphere for a single point.
SphericalHarmonicTransform	class	Dense, differentiable spherical harmonic analysis/synthesis on a grid.
SphericalSpectralConv	class	Spherical spectral convolution: per-degree channel mixing in SH space.
SphericalWaveletConv	class	Per-scale channel mixing in the spherical-wavelet domain.
SphericalWaveletTransform	class	Scale-discretised (needlet) wavelet analysis/synthesis on the sphere.
SqueezeExcitation	class	Squeeze-and-Excitation channel gating (Hu et al., 2018).
Stage	class	One encoder/decoder stage: a leaf block, then an optional nested U-Net.
StandardizedConv	class	Convolution with optional on-the-fly weight standardization.
TTTensor	class	Tensor-Train (Matrix-Product) factorization.
TuckerTensor	class	Tucker factorization: a small core contracted with per-mode factor matrices.
UNet	class	Dimension-flexible U-Net with modern blocks, attention, and nested stages.
Upsample	class	Factor-2 upsample via pixel-shuffle: 1x1 conv, then depth-to-space.
WNO	class	Wavelet Neural Operator: lifting → wavelet blocks → projection.
WNOBlock	class	One wavelet layer: wavelet conv + pointwise skip, then an activation.
WaveletAttention	class	Self-attention performed inside the wavelet domain (Wang et al., 2026).
WaveletConv	class	Wavelet-domain channel mixing — the wavelet analogue of SpectralConv.
WaveletNeuralOperator	class	Wavelet Neural Operator: lifting → wavelet blocks → projection.
WindowedAttention	class	Swin-style windowed multi-head self-attention with cosine similarity.
XUNet	class	Dimension-flexible U-Net with modern blocks, attention, and nested stages.
apply_rank1_proj	function	Apply y_i = ((x ⊙ s_i) @ W) ⊙ r_i + b_i per ensemble member.
band_limited_resize	function	Resample a periodic field onto a new grid by Fourier truncation / zero-padding.
build_siren_specs	function	Produce per-layer specs for a depth-depth SIREN, first + hidden… + last.
cno_receptive_radius	function	Effective receptive radius, in grid points, of a CNO run on its input grid.
hetero_noisy_logits	function	Sample num_mc_samples heteroscedastic logits for one example.
init_factorized_tensor	function	Construct a factorized tensor of the requested kind.
init_rank1_proj	function	Glorot-init shared kernel + per-member rank-1 vectors + zero bias.
log_spaced_omegas	function	Return num_scales log-spaced frequencies spanning the closed band.
mfn_forward	function	Pure-JAX MFN forward pass given user-supplied filter and linear callables.
orthogonal_blocks	function	Build the ORF frequency matrix as [Q_1 S_1, ..., Q_K S_K].
rff_cosine_forward	function	Shared single-cosine RFF feature map: sqrt(2/D) cos(xW/l + b).
rff_forward	function	Shared RFF feature map: sqrt(1/D) [cos(xW/l), sin(xW/l)].
siren_W_limit	function	Return the half-width a of the U(-a, a) weight init for a SIREN layer.

geonnax.basis

Name	Kind	What it does
SlepianBasis	class	Slepian functions of a general region, in the real spherical-harmonic basis.
SlepianCapBasis	class	Precomputed Slepian cap eigendecomposition in the real-SH basis.
associated_legendre	function	Associated Legendre functions \(P_l^m(t)\) for \(0 \le m \le l \le l_{\max}\).
associated_legendre_indices	function	The (l, m) pair of each column of associated_legendre.
cotangent_laplacian_eigpairs	function	Smallest eigenpairs of the cotangent Laplace–Beltrami operator of a mesh.
curlfree_basis	function	Curl-free vector basis: gradients of the box Laplacian eigenfunctions.
divfree_basis	function	Divergence-free vector basis from box-Dirichlet stream functions.
eof_basis	function	Leading n_modes empirical orthogonal functions of a data matrix.
fibonacci_sphere	function	The Fibonacci lattice: n near-uniform points on the unit sphere.
fourier_basis	function	Tensor-product Laplacian eigenpairs on \([-L, L]^D\) (Dirichlet by default).
fourier_basis_1d	function	Evaluate the first num_basis 1D Laplacian eigenfunctions on [-L, L].
fourier_eigenvalues	function	Return the flattened sum-of-squares eigenvalues for a D-dimensional box.
fourier_eigenvalues_1d	function	Return \(\lambda_j = (j\pi / (2L))^2\) for j = 1, ..., num_basis.
fourier_features	function	Cos/sin Fourier basis at dyadic frequencies.
gabor_frame	function	Evaluate a bank of Gabor atoms (Gaussian envelope times oscillation).
gabor_frame_grid	function	Build a dyadic radial-Gabor frame over a bounded box.
gauss_legendre	function	Gauss–Legendre nodes and weights on \([-1, 1]\).
gauss_legendre_grid	function	Gauss–Legendre grid on the sphere with its quadrature weights.
gaussian_window_features	function	Gaussian-window-in-time basis, the temporal analogue of an RBF column.
graph_laplacian_eigpairs	function	Smallest num_basis Laplacian eigenpairs of an undirected graph.
harmonic_degrees	function	Return the per-harmonic degree l for the flattened (l_max + 1)^2 block.
healpix_grid	function	HEALPix pixel centres with their equal-area weights.
icosphere	function	Subdivided icosahedron: a quasi-uniform triangle mesh of the sphere.
interaction_features	function	Element-wise products on selected pairs of input columns.
latlon_grid	function	Regular cell-centred lat/lon grid with exact cell-area weights.
legendre_polynomials	function	Legendre polynomials \(P_0(t), \ldots, P_{l_{\max}}(t)\).
mesh_vertex_areas	function	Lumped (barycentric) vertex areas of a triangle mesh.
needlet_basis	function	Spherical needlets for scales \(j = 0, \ldots, j_{\max}\), plus the constant.
needlet_window	function	The needlet window \(b\): smooth, supported in \([1/B, B]\), a partition of unity.
rbf_basis	function	Evaluate a placeable radial basis at given centres and widths.
real_spherical_harmonics	function	Evaluate real spherical harmonics up to degree l_max.
seasonal_features	function	Cos/sin features at multiples of \(2\pi / \tau_p\).
seasonal_frequencies	function	Flatten (period, harmonic_count) pairs into Python lists.
shannon_number	function	Return the spherical-cap Shannon number for a bandlimit and area.
shell_basis	function	Spherical harmonics times radial polynomials on a spherical shell.
slepian_cap_basis	function	Construct a Slepian basis for an axisymmetric spherical cap.
slepian_cap_eigh_per_m	function	Solve the cap Slepian eigenproblem one azimuthal block at a time.
slepian_concentration_matrix	function	Build the real-SH concentration matrix for a north-pole cap.
slepian_polar_gap_basis	function	Slepian functions of the band between two polar caps (a polar gap).
slepian_region_basis	function	Slepian functions of a region given by a quadrature over it.
spherical_polygon_mask	function	Which points lie inside a spherical polygon.
spherical_rbf_basis	function	Placeable radial basis on the unit 2-sphere, in geodesic distance.
standardize	function	Affine standardize: (x - mu) / std.
unstandardize	function	Inverse of standardize: z * std + mu.
vector_spherical_harmonics	function	Gradient and toroidal vector spherical harmonics for \(1 \le l \le l_{\max}\).
wavelet_basis_1d	function	Orthonormal 1D DWT synthesis-basis matrix on a periodic length-n grid.
wavelet_basis_2d	function	Separable orthonormal 2D DWT basis as a Kronecker product.
wendland_c2	function	Wendland \(C^2\) kernel, positive definite in dimensions \(d \le 3\).
wendland_c4	function	Wendland \(C^4\) kernel, positive definite in dimensions \(d \le 3\).

geonnax.geo

Name	Kind	What it does
EARTH_MEAN_RADIUS_M	constant	Re-exported from builtins.
WGS84_A	constant	Re-exported from builtins.
WGS84_F	constant	Re-exported from builtins.
cartesian_to_tangent	function	3-D vectors to their east/north components (v · ê, v · n̂).
cyclic_encode	function	Encode periodic inputs as concatenated cos/sin features.
deg2rad	function	Convert degrees to radians element-wise.
ecef_to_lonlat	function	Earth-centred, Earth-fixed metres to geodetic (lon, lat) and altitude.
enu_basis	function	Local east/north/up unit vectors at each point.
local_tangent_plane	function	Project points onto the tangent plane at origin (east, north), on a sphere.
lonlat_scale	function	Affine-rescale lon/lat columns.
lonlat_to_cartesian3d	function	Lift lon/lat coordinates onto the unit sphere.
lonlat_to_ecef	function	Geodetic (lon, lat) plus altitude to Earth-centred, Earth-fixed metres.
rotation_to_pole	function	The rotation \(R\) with \(R\,\hat u(\text{centre}) = (0, 0, 1)\).
spherical_harmonic_encode	function	Lift lon/lat to \(S^2\) and evaluate real spherical harmonics.
tangent_to_cartesian	function	East/north components (u, v) to 3-D tangent vectors \(u\hat e + v\hat n\).