Inference

GaussianProcess.condition(dataset) returns a ConditionedPosterior. For a complete example, start with First posterior. The posterior accepts a FunctionalQuery for prediction and sampling.

Prediction methods

Method	Result
predict_mean(query)	Named means
predict_variance(query)	Named marginal variances
predict_marginals(query)	Means and marginal variances together
predict_covariance(query)	Full joint covariance
predict(query)	Means and full joint covariance together
sample(key, query, sample_count=...)	Joint Gaussian draws at the query

The marginal methods avoid constructing a full query covariance. Use predict_marginals for pointwise uncertainty bands. Use a full covariance when correlations between query outputs matter.

Named results and shapes

For a query block named "target" with public shape s:

Access	Shape
prediction.mean["target"]	s
marginals.variance["target"]	s
prediction.covariance.variance["target"]	s
prediction.covariance.block("left", "right")	(*left_shape, *right_shape)
samples["target"]	(sample_count, *s)

A block at n locations has shape (n, *operator_output_shape). For a joint query with a scalar block and a gradient block in d dimensions, their cross-covariance has shape (n, n, d). Sample rows correspond to the same joint draw across blocks. Query predictions describe latent functionals. To obtain a covariance for future noisy readings, add their measurement-error variances to the appropriate diagonal entries.

Prediction results can cross an outer jit. Under vmap, named arrays retain the leading batch axes. Dataset and query input dimensions must agree.

Conditioning and numerical options

gp.condition(dataset, jitter=1e-6) adds jitter to the observation covariance diagonal before solving. Jitter is an absolute covariance quantity and must be finite and nonnegative. Choose it in the units of your model and keep it consistent between fitting and conditioning. Observation noise is specified by the dataset's design. Sampling has its own jitter argument for stabilizing the query covariance factorization.

The default dense solver supports first-order forward and reverse derivatives through conditioning and updates. Higher-order derivatives through its factorization and updates are unsupported. Failed Cholesky factorizations can produce non-finite results, including under JIT. Check outputs before using them. Use consistent floating precision, sensible units, and explicit regularization for redundant measurements. A jitter too small relative to the covariance diagonal can round away.

GaussianProcess(kernel, covariance_tile_size=1024) evaluates observation covariances in smaller tiles. Tile size counts scalar outputs, with at least one location's complete output per tile. Smaller tiles reduce temporary memory at the cost of more loop overhead. The final dense covariance and its factorization still have to fit in memory. The default is None. See numerical backends for solver configuration and interfaces.

Proposed measurements and posterior updates

gp.prepare_conditioning(design) returns a ConditioningState for covariance calculations before outcomes are available. PreparedQuery(conditioning, query) reuses a fixed target when comparing proposed measurements. Preparation stores the full target covariance, including for subsequent variance-only requests. Rebuild it after changing the GP, existing design, jitter, or target.

posterior.append(locations, readings, block="name") updates an existing posterior within a declared block's capacity. posterior.conditioning.append performs a covariance update without readings. See Experimental design and Sequential observations for their use, and observation capacity for compilation and model-reuse conditions.

Classes and methods

gpex.gp

Exact Gaussian-process inference.

GaussianProcess

GaussianProcess(
    kernel: KernelType,
    mean: Mean | ScalarFunction | None = None,
    *,
    covariance_tile_size: int | None = None,
    solver: SolverConfiguration | None = None,
)

Bases: Module

A scalar latent GP whose observations and queries may use linear operators.

Parameters:

Name	Type	Description	Default
kernel	KernelType	Scalar covariance kernel evaluated on points of shape (d,).	required
mean	Mean | ScalarFunction | None	Scalar mean function. None (the default) selects ZeroMean.	None
covariance_tile_size	int | None	Positive tile size for materialized covariance construction. Defaults to None for untiled construction.	None
solver	SolverConfiguration | None	Numerical backend, such as DenseBackend(). None (the default) selects dense Cholesky inference.	None

Source code in src/gpex/gp.py

def __init__(
    self,
    kernel: KernelType,
    mean: Mean | ScalarFunction | None = None,
    *,
    covariance_tile_size: int | None = None,
    solver: SolverConfiguration | None = None,
) -> None:
    if covariance_tile_size is not None and (
        isinstance(covariance_tile_size, bool)
        or not isinstance(covariance_tile_size, int)
        or covariance_tile_size < 1
    ):
        raise ValueError("covariance_tile_size must be a positive integer or None")
    self.kernel = kernel
    self.mean = ZeroMean() if mean is None else mean
    self.covariance_tile_size = covariance_tile_size
    self.solver = DenseBackend() if solver is None else solver

kernel instance-attribute

kernel: KernelType = kernel

mean instance-attribute

mean: Mean | ScalarFunction = (
    ZeroMean() if mean is None else mean
)

covariance_tile_size class-attribute instance-attribute

covariance_tile_size: int | None = covariance_tile_size

solver instance-attribute

solver: SolverConfiguration = (
    DenseBackend() if solver is None else solver
)

prior_mean

prior_mean(query: FunctionalQuery) -> StructuredValues

Return prior means for a query, with named arrays in each block's shape.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and linear operators to evaluate.	required

Returns:

Type	Description
StructuredValues	Structured means with flat shape (q,) for q scalar outputs.

Source code in src/gpex/gp.py

def prior_mean(self, query: FunctionalQuery) -> StructuredValues:
    """Return prior means for a query, with named arrays in each block's shape.

    Args:
        query: Locations and linear operators to evaluate.

    Returns:
        Structured means with ``flat`` shape ``(q,)`` for ``q`` scalar outputs.
    """
    flat = assemble_mean(query.blocks, self.mean)
    return StructuredValues(flat, query.index_map)

prior_covariance

prior_covariance(
    query: FunctionalQuery,
) -> CovarianceResult

Return prior covariance for a query, without observation noise.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and linear operators to evaluate.	required

Returns:

Type	Description
CovarianceResult	Full covariance of shape (q, q), named blocks, and variances.

Source code in src/gpex/gp.py

def prior_covariance(self, query: FunctionalQuery) -> CovarianceResult:
    """Return prior covariance for a query, without observation noise.

    Args:
        query: Locations and linear operators to evaluate.

    Returns:
        Full covariance of shape ``(q, q)``, named blocks, and variances.
    """
    covariance = assemble_covariance(
        query.blocks,
        query.blocks,
        self.kernel,
        tile_size=self.covariance_tile_size,
    )
    covariance = symmetrize(covariance)
    return CovarianceResult(covariance, query.index_map)

prior_variance

prior_variance(query: FunctionalQuery) -> StructuredValues

Return latent marginal variances without forming a full covariance.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and linear operators to evaluate.	required

Returns:

Type	Description
StructuredValues	Variances accessible by block name, in each block's output shape.

Source code in src/gpex/gp.py

def prior_variance(self, query: FunctionalQuery) -> StructuredValues:
    """Return latent marginal variances without forming a full covariance.

    Args:
        query: Locations and linear operators to evaluate.

    Returns:
        Variances accessible by block name, in each block's output shape.
    """
    variance = assemble_covariance_diagonal(query.blocks, self.kernel)
    return StructuredValues(variance, query.index_map)

sample_prior

sample_prior(
    key: PRNGKeyArray,
    query: FunctionalQuery,
    *,
    sample_count: int = 1,
    jitter: float = 1e-06,
) -> StructuredValues

Draw joint samples at a fixed query from the latent prior.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX random key; supply a fresh key for independent draws.	required
query	FunctionalQuery	Locations and linear operators to sample jointly.	required
sample_count	int	Number of draws. Defaults to 1.	1
jitter	float	Diagonal stabilization for sampling. Defaults to 1e-6.	1e-06

Returns:

Type	Description
StructuredValues	Samples with flat shape (sample_count, q) and named arrays of shape (sample_count, *block.shape).

Source code in src/gpex/gp.py

def sample_prior(
    self,
    key: PRNGKeyArray,
    query: FunctionalQuery,
    *,
    sample_count: int = 1,
    jitter: float = 1e-6,
) -> StructuredValues:
    """Draw joint samples at a fixed query from the latent prior.

    Args:
        key: JAX random key; supply a fresh key for independent draws.
        query: Locations and linear operators to sample jointly.
        sample_count: Number of draws. Defaults to 1.
        jitter: Diagonal stabilization for sampling. Defaults to ``1e-6``.

    Returns:
        Samples with ``flat`` shape ``(sample_count, q)`` and named arrays of shape
            ``(sample_count, *block.shape)``.
    """
    mean = self.prior_mean(query)
    covariance = self.prior_covariance(query)
    samples = _sample_gaussian(
        key, mean.flat, covariance.full, sample_count, jitter
    )
    return StructuredValues(samples, query.index_map)

prepare_conditioning

prepare_conditioning(
    design: ObservationDesign | FunctionalDataset,
    *,
    jitter: float = 1e-06,
) -> ConditioningState[KernelType]

Prepare conditional covariance without using observation outcomes.

Parameters:

Name	Type	Description	Default
design	ObservationDesign | FunctionalDataset	Measurement locations, operators, and noise. A dataset is also accepted; its observed values are ignored.	required
jitter	float	Diagonal stabilization of the observation covariance. Defaults to 1e-6.	1e-06

Returns:

Type	Description
ConditioningState[KernelType]	Reusable conditioning state for solves, covariance prediction, and
ConditioningState[KernelType]	design updates.

Source code in src/gpex/gp.py

def prepare_conditioning(
    self,
    design: ObservationDesign | FunctionalDataset,
    *,
    jitter: float = 1e-6,
) -> ConditioningState[KernelType]:
    """Prepare conditional covariance without using observation outcomes.

    Args:
        design: Measurement locations, operators, and noise. A dataset is
            also accepted; its observed values are ignored.
        jitter: Diagonal stabilization of the observation covariance.
            Defaults to ``1e-6``.

    Returns:
        Reusable conditioning state for solves, covariance prediction, and
        design updates.
    """
    if isinstance(design, FunctionalDataset):
        design = design.design
    measurements = prepare_functionals(design.blocks)
    system = CovarianceSystem(
        measurements,
        self.kernel,
        design.variance_flat,
        jitter,
    )
    linear_system = self.solver.prepare(
        system,
        materialization_tile_size=self.covariance_tile_size,
        active=measurements.active,
    )
    return ConditioningState(
        self,
        system,
        linear_system,
        design.index_map,
        measurement_block_layout(design.blocks),
    )

condition

condition(
    dataset: FunctionalDataset, *, jitter: float = 1e-06
) -> ConditionedPosterior[KernelType]

Condition on a dataset with gp.condition(dataset, jitter=1e-6).

The callable captures this immutable GP and can be compiled directly: condition = eqx.filter_jit(gp.condition). Reuse it across datasets with the same reserved block structure and dtypes. Values, locations, noise and active counts may change without compiling again. To change or differentiate model parameters, pass the GP into a compiled function.

Returns a posterior for the current contents on every call. This does not reuse a previous factorization. Dense work scales with the reserved scalar capacity, including unused slots.

Parameters:

Name	Type	Description	Default
dataset	FunctionalDataset	Named observation blocks with reserved locations, observed operator outputs, active counts and scalar noise variances.	required
jitter	float	Diagonal stabilization of the observation covariance. Defaults to 1e-6. Unused slots retain a unit diagonal.	1e-06

Returns:

Type	Description
ConditionedPosterior[KernelType]	A posterior supporting prediction and sampling. Failed dense factorization propagates non-finite numerical results.

Source code in src/gpex/gp.py

def condition(
    self, dataset: FunctionalDataset, *, jitter: float = 1e-6
) -> ConditionedPosterior[KernelType]:
    """Condition on a dataset with ``gp.condition(dataset, jitter=1e-6)``.

    The callable captures this immutable GP and can be compiled directly:
    ``condition = eqx.filter_jit(gp.condition)``. Reuse it across datasets
    with the same reserved block structure and dtypes. Values, locations,
    noise and active counts may change without compiling again. To change
    or differentiate model parameters, pass the GP into a compiled function.

    Returns a posterior for the current contents on every call. This does
    not reuse a previous factorization. Dense work scales with the reserved
    scalar capacity, including unused slots.

    Args:
        dataset: Named observation blocks with reserved locations, observed
            operator outputs, active counts and scalar noise variances.
        jitter: Diagonal stabilization of the observation covariance.
            Defaults to ``1e-6``. Unused slots retain a unit diagonal.

    Returns:
        A posterior supporting prediction and sampling. Failed dense
            factorization propagates non-finite numerical results.
    """
    conditioning = self.prepare_conditioning(dataset, jitter=jitter)
    observation_mean = assemble_mean(conditioning.measurements, self.mean)
    values = dataset.values_flat
    alpha = conditioning.solve(values - observation_mean)
    return ConditionedPosterior(conditioning, values, alpha)

ConditioningState

ConditioningState(
    gp: GaussianProcess[KernelType],
    system: CovarianceSystem,
    linear_system: LinearSystem,
    index_map: IndexMap,
    block_layout: tuple[tuple[int, int], ...],
)

Bases: Module

Outcome-independent covariance state created by prepare_conditioning.

Use GaussianProcess.prepare_conditioning to construct this state from a design or dataset. Its solves act on flattened scalar observation outputs.

Source code in src/gpex/gp.py

def __init__(
    self,
    gp: GaussianProcess[KernelType],
    system: CovarianceSystem,
    linear_system: LinearSystem,
    index_map: IndexMap,
    block_layout: tuple[tuple[int, int], ...],
) -> None:
    self.gp = gp
    self.measurements = system.measurements
    self.variance_flat = system.variance
    self.index_map = index_map
    self.block_layout = block_layout
    self.linear_system = linear_system
    self.jitter = system.jitter

gp instance-attribute

gp: GaussianProcess[KernelType] = gp

measurements instance-attribute

measurements: PreparedFunctionals = system.measurements

variance_flat instance-attribute

variance_flat: ObservationVector = system.variance

index_map class-attribute instance-attribute

index_map: IndexMap = index_map

block_layout class-attribute instance-attribute

block_layout: tuple[tuple[int, int], ...] = block_layout

linear_system instance-attribute

linear_system: LinearSystem = linear_system

jitter class-attribute instance-attribute

jitter: float = system.jitter

solve

solve(
    right_hand_side: FunctionalVector,
) -> FunctionalVector

Solve the stabilized observation covariance system.

Parameters:

Name	Type	Description	Default
right_hand_side	FunctionalVector	Array with shape (N, *rhs_shape), where N is the number of scalar observation outputs.	required

Returns:

Type	Description
FunctionalVector	Solution with the same shape as right_hand_side.

Source code in src/gpex/gp.py

def solve(
    self,
    right_hand_side: FunctionalVector,
) -> FunctionalVector:
    """Solve the stabilized observation covariance system.

    Args:
        right_hand_side: Array with shape ``(N, *rhs_shape)``, where ``N`` is
            the number of scalar observation outputs.

    Returns:
        Solution with the same shape as ``right_hand_side``.
    """
    return self.linear_system.solve(right_hand_side)

predict_covariance

predict_covariance(
    query: FunctionalQuery,
) -> CovarianceResult

Return conditional latent covariance at a query, without outcome values.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and operators with the training input dimension.	required

Returns:

Type	Description
CovarianceResult	Full covariance of shape (q, q), named blocks, and variances. Observation noise is not added to the query covariance.

Source code in src/gpex/gp.py

def predict_covariance(self, query: FunctionalQuery) -> CovarianceResult:
    """Return conditional latent covariance at a query, without outcome values.

    Args:
        query: Locations and operators with the training input dimension.

    Returns:
        Full covariance of shape ``(q, q)``, named blocks, and variances.
            Observation noise is not added to the query covariance.
    """
    prior = assemble_covariance(
        query.blocks,
        query.blocks,
        self.gp.kernel,
        tile_size=self.gp.covariance_tile_size,
    )
    cross = _assemble_query_cross(self, query)
    covariance = symmetrize(prior - self.linear_system.covariance_reduction(cross))
    return CovarianceResult(covariance, query.index_map)

predict_variance

predict_variance(
    query: FunctionalQuery,
) -> StructuredValues

Return conditional latent variances without a full query covariance.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and operators with the training input dimension.	required

Returns:

Type	Description
StructuredValues	Variances accessible by block name, without observation noise. Storage still includes the observation-query cross-covariance.

Source code in src/gpex/gp.py

def predict_variance(self, query: FunctionalQuery) -> StructuredValues:
    """Return conditional latent variances without a full query covariance.

    Args:
        query: Locations and operators with the training input dimension.

    Returns:
        Variances accessible by block name, without observation noise.
            Storage still includes the observation-query cross-covariance.
    """
    prior = assemble_covariance_diagonal(query.blocks, self.gp.kernel)
    cross = _assemble_query_cross(self, query)
    variance = prior - self.linear_system.variance_reduction(cross)
    return StructuredValues(variance, query.index_map)

append

append(
    locations: Locations, *, block: str
) -> ConditioningState[KernelType]

Activate locations in a declared block, retaining its operator/noise.

Existing measurements and numerical settings stay fixed. Shapes and index maps retain their reserved sizes. Unknown names and capacity overflow raise, including overflow under JIT.

Source code in src/gpex/gp.py

def append(
    self, locations: Locations, *, block: str
) -> ConditioningState[KernelType]:
    """Activate locations in a declared block, retaining its operator/noise.

    Existing measurements and numerical settings stay fixed. Shapes and
    index maps retain their reserved sizes. Unknown names and capacity
    overflow raise, including overflow under JIT.
    """
    update = self._prepare_update(locations, block=block)
    linear_system = self.linear_system.append(update)
    return ConditioningState(
        self.gp, update.system, linear_system, self.index_map, self.block_layout
    )

append_solve

append_solve(
    locations: Locations,
    solution: ObservationVector,
    right_hand_side: ObservationVector,
    *,
    block: str,
) -> tuple[
    ConditioningState[KernelType], ObservationVector
]

Activate locations and update a canonical-order vector solution.

solution and the new RHS both have the reserved scalar size. The RHS must retain all existing active entries and be zero in unused slots. Dense updates reuse the factor and update the existing coefficients.

Source code in src/gpex/gp.py

def append_solve(
    self,
    locations: Locations,
    solution: ObservationVector,
    right_hand_side: ObservationVector,
    *,
    block: str,
) -> tuple[ConditioningState[KernelType], ObservationVector]:
    """Activate locations and update a canonical-order vector solution.

    ``solution`` and the new RHS both have the reserved scalar size. The
    RHS must retain all existing active entries and be zero in unused slots.
    Dense updates reuse the factor and update the existing coefficients.
    """
    update = self._prepare_update(locations, block=block)
    linear_system, updated_solution = self.linear_system.append_solve(
        update, solution, right_hand_side
    )
    conditioning = ConditioningState(
        self.gp, update.system, linear_system, self.index_map, self.block_layout
    )
    return conditioning, updated_solution

ConditionedPosterior

Bases: Module

An immutable posterior state supporting streaming updates.

conditioning instance-attribute

conditioning: ConditioningState[KernelType]

y_flat instance-attribute

y_flat: ObservationVector

alpha instance-attribute

alpha: ObservationVector

observation_size property

observation_size: int

Number of flattened scalar observation functionals in the state.

predict_mean

predict_mean(query: FunctionalQuery) -> StructuredValues

Return posterior means in each query block's output shape.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Locations and operators with the training input dimension.	required

Returns:

Type	Description
StructuredValues	Structured means with flat shape (q,) and named block arrays.

Source code in src/gpex/gp.py

def predict_mean(self, query: FunctionalQuery) -> StructuredValues:
    """Return posterior means in each query block's output shape.

    Args:
        query: Locations and operators with the training input dimension.

    Returns:
        Structured means with ``flat`` shape ``(q,)`` and named block arrays.
    """
    prior_mean = assemble_mean(query.blocks, self.conditioning.gp.mean)
    cross = _assemble_query_cross(self.conditioning, query)
    flat = prior_mean + jnp.matmul(
        cross.T, self.alpha, precision=jax.lax.Precision.HIGHEST
    )
    return StructuredValues(flat, query.index_map)

predict_covariance

predict_covariance(
    query: FunctionalQuery,
) -> CovarianceResult

Delegate to predict_covariance.

Source code in src/gpex/gp.py

def predict_covariance(self, query: FunctionalQuery) -> CovarianceResult:
    """Delegate to
    [`predict_covariance`][gpex.gp.ConditioningState.predict_covariance].
    """
    return self.conditioning.predict_covariance(query)

predict_variance

predict_variance(
    query: FunctionalQuery,
) -> StructuredValues

Delegate to predict_variance.

Source code in src/gpex/gp.py

def predict_variance(self, query: FunctionalQuery) -> StructuredValues:
    """Delegate to
    [`predict_variance`][gpex.gp.ConditioningState.predict_variance].
    """
    return self.conditioning.predict_variance(query)

predict_marginals

predict_marginals(
    query: FunctionalQuery,
) -> MarginalPrediction

Return latent means and marginal variances, sharing cross-covariance work.

The result contains named arrays in each query block's output shape and can cross an outer jit or vmap boundary. Storage includes the observation/query cross-covariance, but no full query covariance. Observation noise is excluded.

Source code in src/gpex/gp.py

def predict_marginals(self, query: FunctionalQuery) -> MarginalPrediction:
    """Return latent means and marginal variances, sharing cross-covariance work.

    The result contains named arrays in each query block's output shape and
    can cross an outer ``jit`` or ``vmap`` boundary. Storage includes the
    observation/query cross-covariance, but no full query covariance.
    Observation noise is excluded.
    """
    conditioning = self.conditioning
    cross = _assemble_query_cross(conditioning, query)
    mean = assemble_mean(query.blocks, conditioning.gp.mean) + jnp.matmul(
        cross.T, self.alpha, precision=jax.lax.Precision.HIGHEST
    )
    variance = assemble_covariance_diagonal(query.blocks, conditioning.gp.kernel)
    variance = variance - conditioning.linear_system.variance_reduction(cross)
    return MarginalPrediction(
        StructuredValues(mean, query.index_map),
        StructuredValues(variance, query.index_map),
    )

predict

predict(query: FunctionalQuery) -> Prediction

Return the joint latent mean and covariance at a query.

Parameters:

Name	Type	Description	Default
query	FunctionalQuery	Named blocks of locations and linear operators to predict.	required

Returns:

Type	Description
Prediction	A Prediction with structured means and a full (q, q) covariance for q flattened query outputs. The complete result is a PyTree and can be returned from JIT-compiled functions.

Source code in src/gpex/gp.py

def predict(self, query: FunctionalQuery) -> Prediction:
    """Return the joint latent mean and covariance at a query.

    Args:
        query: Named blocks of locations and linear operators to predict.

    Returns:
        A ``Prediction`` with structured means and a full ``(q, q)`` covariance for
            ``q`` flattened query outputs. The complete result is a PyTree and can
            be returned from JIT-compiled functions.
    """
    conditioning = self.conditioning
    gp = conditioning.gp
    cross = _assemble_query_cross(conditioning, query)
    prior_mean = assemble_mean(query.blocks, gp.mean)
    mean = StructuredValues(
        prior_mean
        + jnp.matmul(cross.T, self.alpha, precision=jax.lax.Precision.HIGHEST),
        query.index_map,
    )
    prior_covariance = assemble_covariance(
        query.blocks,
        query.blocks,
        gp.kernel,
        tile_size=gp.covariance_tile_size,
    )
    covariance = symmetrize(
        prior_covariance - conditioning.linear_system.covariance_reduction(cross)
    )
    return Prediction(mean, CovarianceResult(covariance, query.index_map))

sample

sample(
    key: PRNGKeyArray,
    query: FunctionalQuery,
    *,
    sample_count: int = 1,
    jitter: float = 1e-06,
) -> StructuredValues

Draw joint latent posterior samples at a fixed query.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX random key; supply a fresh key for independent draws.	required
query	FunctionalQuery	Locations and linear operators to sample jointly.	required
sample_count	int	Number of draws. Defaults to 1.	1
jitter	float	Diagonal stabilization for sampling. Defaults to 1e-6.	1e-06

Returns:

Type	Description
StructuredValues	Samples with flat shape (sample_count, q) and named arrays of shape (sample_count, *block.shape). Observation noise is excluded. Use FeaturePosteriorPath(posterior, key) for a callable approximate sample path.

Source code in src/gpex/gp.py

def sample(
    self,
    key: PRNGKeyArray,
    query: FunctionalQuery,
    *,
    sample_count: int = 1,
    jitter: float = 1e-6,
) -> StructuredValues:
    """Draw joint latent posterior samples at a fixed query.

    Args:
        key: JAX random key; supply a fresh key for independent draws.
        query: Locations and linear operators to sample jointly.
        sample_count: Number of draws. Defaults to 1.
        jitter: Diagonal stabilization for sampling. Defaults to ``1e-6``.

    Returns:
        Samples with ``flat`` shape ``(sample_count, q)`` and named arrays of shape
            ``(sample_count, *block.shape)``. Observation noise is excluded. Use
            ``FeaturePosteriorPath(posterior, key)`` for a callable approximate
            sample path.
    """
    prediction = self.predict(query)
    samples = _sample_gaussian(
        key, prediction.mean.flat, prediction.covariance.full, sample_count, jitter
    )
    return StructuredValues(samples, query.index_map)

append

append(
    locations: Locations,
    values: FunctionalValues,
    *,
    block: str,
) -> ConditionedPosterior[KernelType]

Add observations to a declared block without enlarging storage.

locations has shape (m, d). values has shape (m, *operator_output_shape) for the selected block. Its operator, noise variance, model parameters and previous observations remain fixed. Returns a new posterior; overflow raises instead of resizing.

Repeated fixed-shape batches into the same block reuse an outer JIT. Dense systems update their factor and coefficients through the Schur complement.

Source code in src/gpex/gp.py

def append(
    self,
    locations: Locations,
    values: FunctionalValues,
    *,
    block: str,
) -> ConditionedPosterior[KernelType]:
    """Add observations to a declared block without enlarging storage.

    ``locations`` has shape ``(m, d)``. ``values`` has shape
    ``(m, *operator_output_shape)`` for the selected block. Its operator,
    noise variance, model parameters and previous observations remain fixed.
    Returns a new posterior; overflow raises instead of resizing.

    Repeated fixed-shape batches into the same block reuse an outer JIT.
    Dense systems update their factor and coefficients through the Schur
    complement.
    """
    index = self.conditioning.index_map.block(block)
    values = jnp.asarray(values)
    expected = (locations.count, *index.shape[1:])
    if values.shape != expected:
        raise ValueError(
            f"appended values have shape {values.shape}; expected {expected}"
        )
    update = self.conditioning._prepare_update(locations, block=block)
    y_flat = self.y_flat.at[update.indices].set(values.reshape(-1))
    right_hand_side = y_flat - assemble_mean(
        update.system.measurements, self.conditioning.gp.mean
    )
    linear_system, alpha = self.conditioning.linear_system.append_solve(
        update, self.alpha, right_hand_side
    )
    conditioning = ConditioningState(
        self.conditioning.gp,
        update.system,
        linear_system,
        self.conditioning.index_map,
        self.conditioning.block_layout,
    )
    return ConditionedPosterior(conditioning, y_flat, alpha)

gpex.results

Structured latent predictions and covariance results.

Prediction

Bases: Module

Joint latent predictions returned by ConditionedPosterior.predict.

Attributes:

Name	Type	Description
mean	StructuredValues	Means accessible as mean[block_name], in each block's shape.
covariance	CovarianceResult	Full covariance, named cross-blocks, and marginal variances.

mean instance-attribute

mean: StructuredValues

covariance instance-attribute

covariance: CovarianceResult

MarginalPrediction

Bases: Module

Latent means and variances without a full query covariance.

Both fields contain flat arrays of shape (q,) and named views matching each query block. Leading batch axes from vmap are preserved.

mean instance-attribute

mean: StructuredValues

variance instance-attribute

variance: StructuredValues

CovarianceResult

CovarianceResult(
    full: CovarianceMatrix, index_map: IndexMap
)

Bases: Module

Covariance of the latent query outputs, without observation noise.

Parameters:

Name	Type	Description	Default
full	CovarianceMatrix	Dense covariance with shape (q, q) for q scalar outputs.	required
index_map	IndexMap	Static map between named blocks and flattened output indices.	required

Attributes:

Name	Type	Description
full	CovarianceMatrix	Dense latent covariance.
variance	StructuredValues	Diagonal entries, accessible by block name in each block's shape.
index_map	IndexMap	Mapping used for named covariance blocks and variances.

Source code in src/gpex/results.py

def __init__(self, full: CovarianceMatrix, index_map: IndexMap) -> None:
    full = jnp.asarray(full)
    expected_shape = (index_map.size, index_map.size)
    if full.shape != expected_shape:
        raise ValueError(
            f"covariance has shape {full.shape}; expected {expected_shape}"
        )
    self.full = full
    self.index_map = index_map

full instance-attribute

full: CovarianceMatrix = full

index_map class-attribute instance-attribute

index_map: IndexMap = index_map

variance property

variance: StructuredValues

Marginal variances derived from the full covariance.

block

block(left_name: str, right_name: str) -> CovarianceBlock

Return a named cross-covariance with shape (*left.shape, *right.shape).

Parameters:

Name	Type	Description	Default
left_name	str	Row block name from the query.	required
right_name	str	Column block name from the query.	required

Returns:

Type	Description
CovarianceBlock	The covariance between the two blocks, preserving their output axes. Leading batch axes introduced by vmap are preserved too.

Source code in src/gpex/results.py

def block(self, left_name: str, right_name: str) -> CovarianceBlock:
    """Return a named cross-covariance with shape ``(*left.shape, *right.shape)``.

    Args:
        left_name: Row block name from the query.
        right_name: Column block name from the query.

    Returns:
        The covariance between the two blocks, preserving their output axes.
            Leading batch axes introduced by ``vmap`` are preserved too.
    """
    left = self.index_map.block(left_name)
    right = self.index_map.block(right_name)
    leading_shape = self.full.shape[:-2]
    return self.full[..., left.slice, right.slice].reshape(
        (*leading_shape, *left.shape, *right.shape)
    )

gpex.prepared_query.PreparedQuery

PreparedQuery(
    conditioning: ConditioningState, query: FunctionalQuery
)

Bases: Module

A fixed target and conditioning state cached for comparing measurements.

Construct with PreparedQuery(conditioning, query). Each *_after call evaluates one proposal against the original conditioning state; calls do not accumulate measurements. Prepare a new query after changing the GP, observation design, or target. The cache is an immutable PyTree.

Candidate noise and the conditioning's jitter enter the small conditional measurement system. Target covariance remains latent. Storage is O(N*q + q*q) for N observations and q target outputs.

Parameters:

Name	Type	Description	Default
conditioning	ConditioningState	Original conditioning state.	required
query	FunctionalQuery	Fixed target locations and operators.	required

Attributes:

Name	Type	Description
conditioning	ConditioningState	Original conditioning state.
covariance	CovarianceResult	Target covariance before a proposed measurement.
projection	CovarianceProjection	Cached backend factors for observation-query reductions.

Source code in src/gpex/prepared_query.py

def __init__(self, conditioning: ConditioningState, query: FunctionalQuery) -> None:
    target = prepare_functionals(query.blocks)
    prior = assemble_covariance(
        target,
        target,
        conditioning.gp.kernel,
        tile_size=conditioning.gp.covariance_tile_size,
    )
    cross = assemble_covariance(
        conditioning.measurements,
        target,
        conditioning.gp.kernel,
        tile_size=conditioning.gp.covariance_tile_size,
    )
    projection = conditioning.linear_system.project(cross)
    self.conditioning = conditioning
    self.target = target
    self.covariance = CovarianceResult(
        symmetrize(prior - projection.gram()), query.index_map
    )
    self.projection = projection

conditioning instance-attribute

conditioning: ConditioningState = conditioning

target instance-attribute

target: PreparedFunctionals = target

covariance instance-attribute

covariance: CovarianceResult = CovarianceResult(
    symmetrize(prior - projection.gram()), query.index_map
)

projection instance-attribute

projection: CovarianceProjection = projection

predict_covariance_after

predict_covariance_after(
    block: ObservationDesignBlock,
) -> CovarianceResult

Return target covariance after assimilating a proposed measurement.

Parameters:

Name	Type	Description	Default
block	ObservationDesignBlock	Proposed measurement design, including its shared noise variance. For observed data, pass observation.design.	required

Returns:

Type	Description
CovarianceResult	Full latent target covariance, equivalent to appending the block and predicting at the target, up to numerical solver accuracy. Candidate coordinates, operator parameters, and noise can be differentiated through this operation.

Source code in src/gpex/prepared_query.py

def predict_covariance_after(
    self,
    block: ObservationDesignBlock,
) -> CovarianceResult:
    """Return target covariance after assimilating a proposed measurement.

    Args:
        block: Proposed measurement design, including its shared noise
            variance. For observed data, pass ``observation.design``.

    Returns:
        Full latent target covariance, equivalent to appending the block
            and predicting at the target, up to numerical solver accuracy.
            Candidate coordinates, operator parameters, and noise can be
            differentiated through this operation.
    """
    system, cross = self._conditional_measurement(block)
    covariance = self.covariance.full - system.covariance_reduction(cross)
    return CovarianceResult(symmetrize(covariance), self.covariance.index_map)

predict_variance_after

predict_variance_after(
    block: ObservationDesignBlock,
) -> StructuredValues

Return target variances after a proposal without a full update matrix.

Parameters:

Name	Type	Description	Default
block	ObservationDesignBlock	Proposed measurement design, including its shared noise variance. For observed data, pass observation.design.	required

Returns:

Type	Description
StructuredValues	Differentiable latent variances by target block name. Initial query preparation still stores the current full target covariance.

Source code in src/gpex/prepared_query.py

def predict_variance_after(
    self,
    block: ObservationDesignBlock,
) -> StructuredValues:
    """Return target variances after a proposal without a full update matrix.

    Args:
        block: Proposed measurement design, including its shared noise
            variance. For observed data, pass ``observation.design``.

    Returns:
        Differentiable latent variances by target block name. Initial query
            preparation still stores the current full target covariance.
    """
    system, cross = self._conditional_measurement(block)
    variance = self.covariance.variance.flat - system.variance_reduction(cross)
    return StructuredValues(variance, self.covariance.index_map)

Feature-based callable draws require a FeatureKernel. RBF and Matérn families provide Fourier features, while custom feature kernels can use other bases. GaussianProcess, ConditioningState, and ConditionedPosterior preserve the kernel type so this requirement is visible to static type checkers.

gpex.feature_paths

Callable posterior draws using kernel-provided finite features.

FeaturePosteriorPath

FeaturePosteriorPath(
    posterior: ConditionedPosterior[KernelType],
    key: PRNGKeyArray,
    *,
    feature_count: int = 1024,
)

Bases: Module

A callable posterior draw built from a finite kernel feature map.

The path uses the kernel's random features for the prior and a pathwise conditioning correction for the observations. Repeated calls evaluate the same draw, so the object can be vectorized, differentiated, or passed to a LinearOperator. Construct with FeaturePosteriorPath(posterior, key). For exact samples at a fixed finite query, use posterior.sample(key, query). Differentiating a finite feature expansion beyond the GP's smoothness does not define a valid GP derivative.

Parameters:

Name	Type	Description	Default
posterior	ConditionedPosterior[KernelType]	Conditioned GP whose kernel implements FeatureKernel. RBF and Matérn families provide Fourier features. Custom feature kernels may provide other finite bases.	required
key	PRNGKeyArray	JAX random key for this draw.	required
feature_count	int	Number of features. Defaults to 1024. For the built-in kernels, larger maps improve the prior approximation statistically.	1024

Source code in src/gpex/feature_paths.py

def __init__(
    self,
    posterior: ConditionedPosterior[KernelType],
    key: PRNGKeyArray,
    *,
    feature_count: int = 1024,
) -> None:
    conditioning = posterior.conditioning
    kernel = conditioning.gp.kernel
    if not isinstance(kernel, FeatureKernel):
        raise TypeError(
            "FeaturePosteriorPath requires a FeatureKernel; "
            f"got {type(kernel).__name__}. "
            "Use posterior.sample(key, query) for finite-query samples."
        )
    input_dimension = conditioning.measurements.input_dimension
    feature_key, weight_key, noise_key = jax.random.split(key, 3)
    features = kernel.random_features(
        feature_key, input_dimension, feature_count=feature_count
    )
    prior_path = _FeaturePath(
        features=features,
        weights=jax.random.normal(weight_key, (feature_count,)),
    )

    self.prior_path = prior_path
    self.mean = conditioning.gp.mean
    self.kernel = kernel
    self.measurements = conditioning.measurements
    self.correction_weights = _conditioning_weights(
        posterior, prior_path, noise_key
    )
    self.covariance_tile_size = conditioning.gp.covariance_tile_size
    self.input_dimension = input_dimension

prior_path instance-attribute

prior_path: _FeaturePath = prior_path

mean instance-attribute

mean: ScalarFunction = conditioning.gp.mean

kernel instance-attribute

kernel: KernelType = kernel

measurements instance-attribute

measurements: PreparedFunctionals = (
    conditioning.measurements
)

correction_weights instance-attribute

correction_weights: Array = _conditioning_weights(
    posterior, prior_path, noise_key
)

covariance_tile_size class-attribute instance-attribute

covariance_tile_size: int | None = (
    conditioning.gp.covariance_tile_size
)

input_dimension class-attribute instance-attribute

input_dimension: int = input_dimension