Skip to content

Numerical backends

GaussianProcess accepts a numerical backend through solver. linalg.DenseBackend() is the stateless default backend. Its prepare method returns a CholeskySystem containing the factorization and activity state for exact dense inference and fitting.

Backend contract

SolverConfiguration.prepare(system, ...) prepares an immutable LinearSystem. The input SymmetricLinearOperator supplies capacity, matmul, diagonal, and tiled materialize. Prepared systems provide solves, quadratic forms, log determinants, covariance and variance reductions, reusable projections, and append operations. These interfaces use numerical arrays and activity information rather than GP measurement types.

For a prepared system, size is the active scalar count as a dynamic JAX integer, and capacity is the reserved scalar dimension as a static Python integer. Solves and projections accept full-capacity inputs. Activation increases size while preserving capacity. Input operators also expose capacity, and updates expose previous_capacity for the reserved dimension before activation.

Custom backends inherit from eqx.Module and SolverConfiguration. Their prepared systems inherit from eqx.Module and LinearSystem. Both interfaces are abstract base classes, so subclasses must implement all abstract methods before they can be instantiated. Numerical inputs and update descriptions explicitly inherit the SymmetricLinearOperator and LinearSystemUpdate abstract base classes. Updates store their numerical system and newly active indices as fields.

The small input interface keeps backends independent of kernels and observation types. Dense preparation requests materialization. The operator also supports products without storing the full matrix, ready for future iterative backends.

CovarianceProjection is an abstract interface for reusable covariance reductions. The dense CholeskyProjection stores a whitened cross-covariance whose Gram matrix is the reduction. Diagonal and cross reductions let consumers reuse that state without accessing a backend's factorization.

Likelihood objectives use the configured backend's quadratic_form and logdet. The dense quadratic form needs one triangular solve. Invalid factorizations propagate non-finite statistics so fitting can reject invalid trials. Backend configuration, prepared state, and returned results must support JAX transformations, with numerical arrays stored as dynamic PyTree leaves.

Use configuration.prepare(system) to construct prepared state. For an already materialized matrix, use CholeskySystem.from_matrix(matrix, jitter=...). Factorization uses the symmetric part of the matrix with diagonal jitter. cholops owns this input symmetrization and the symmetry-preserving reverse rule.

Dense systems store a cholops PaddedCholesky in factor. cholops owns factorization, numerical consumers, reserved-slot activation, and coefficient reuse. factor.order maps physical factor rows to canonical slots, and factor.inverse_order maps canonical slots back to physical rows. factor.factor exposes the lower-triangular array in physical order. solve and project accept canonical-order inputs. Preparing a system with active=mask reserves the inactive identity slots for append. Solves, quadratic forms, and projections ignore inactive input entries, including non-finite values. Solves return zero inactive coefficients. Generic systems without a mask are fully active. CholeskySystem(factor, tile_size) wraps an existing PaddedCholesky.

gpex.linalg.DenseBackend

Bases: Module, SolverConfiguration

Stateless backend preparing dense Cholesky systems.

Jitter is owned by the input system. Preparation returns a reusable :class:CholeskySystem containing the factorization and activity state.

prepare 

prepare(
    system: SymmetricLinearOperator,
    *,
    materialization_tile_size: int | None = None,
    active: Activity | None = None,
) -> CholeskySystem

Prepare a factorization from a backend-neutral operator.

If supplied, active identifies used rows. Inactive rows must be an independent identity system. They become reusable factor storage. Numerical consumers ignore inactive input entries. If no activity mask is supplied, every row is active.

Source code in src/gpex/linalg/dense.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def prepare(
    self,
    system: SymmetricLinearOperator,
    *,
    materialization_tile_size: int | None = None,
    active: Activity | None = None,
) -> CholeskySystem:
    """Prepare a factorization from a backend-neutral operator.

    If supplied, ``active`` identifies used rows. Inactive rows must be an
    independent identity system. They become reusable factor storage.
    Numerical consumers ignore inactive input entries. If no activity
    mask is supplied, every row is active.
    """
    matrix = system.materialize(tile_size=materialization_tile_size)
    if active is None:
        active = jnp.ones(system.capacity, dtype=bool)
    factor = PaddedCholesky.from_matrix(matrix, active=active)
    return CholeskySystem(factor, materialization_tile_size)

gpex.linalg.CholeskySystem

Bases: Module, LinearSystem

A dense system backed by a cholops fixed-capacity factor.

cholops owns storage, ordering, inactive-slot masking, and numerical updates. Failed factorizations retain NaNs. Numerical operations support JIT, vmap, and first-order differentiation.

factor instance-attribute 

factor: PaddedCholesky

tile_size class-attribute instance-attribute 

tile_size: int | None = eqx.field(static=True)

size property 

size: Count

Number of active scalar entries, as a dynamic JAX integer.

capacity property 

capacity: int

Reserved scalar dimension, including inactive slots.

from_matrix classmethod 

from_matrix(
    matrix: SquareMatrix, jitter: float
) -> CholeskySystem

Factorize the symmetric part of a square matrix with diagonal jitter.

Source code in src/gpex/linalg/dense.py 
78
79
80
81
82
83
84
85
86
87
88
89
@classmethod
def from_matrix(cls, matrix: SquareMatrix, jitter: float) -> CholeskySystem:
    """Factorize the symmetric part of a square matrix with diagonal jitter."""
    validate_jitter(jitter)
    matrix = jnp.asarray(matrix)
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        raise ValueError("matrix must be square")
    stabilized = matrix + jitter * jnp.eye(matrix.shape[0])
    factor = PaddedCholesky.from_matrix(
        stabilized, active=jnp.ones(matrix.shape[0], dtype=bool)
    )
    return cls(factor, None)

solve 

solve(right_hand_side: RightHandSide) -> RightHandSide

Solve on active rows for vector or array RHS, including empty batches.

Inactive input entries are ignored and inactive solution entries are zero.

Source code in src/gpex/linalg/dense.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def solve(
    self,
    right_hand_side: RightHandSide,
) -> RightHandSide:
    """Solve on active rows for vector or array RHS, including empty batches.

    Inactive input entries are ignored and inactive solution entries are zero.
    """
    right_hand_side = jnp.asarray(right_hand_side)
    if right_hand_side.ndim == 0 or right_hand_side.shape[0] != self.capacity:
        raise ValueError("right-hand side must start with the system capacity")
    shape = right_hand_side.shape
    columns = right_hand_side.reshape((self.capacity, -1))
    return self.factor.solve(columns).reshape(shape)

quadratic_form 

quadratic_form(right_hand_side: Vector) -> Scalar

Return the active inverse quadratic form with one triangular solve.

The vector has shape (capacity,) in canonical order. Supports JIT, vectorization, and differentiation through the factor and vector. Inactive entries are ignored. Failed factorizations propagate non-finite values on the active system.

Source code in src/gpex/linalg/dense.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def quadratic_form(self, right_hand_side: Vector) -> Scalar:
    """Return the active inverse quadratic form with one triangular solve.

    The vector has shape ``(capacity,)`` in canonical order. Supports JIT,
    vectorization, and differentiation through the factor and vector.
    Inactive entries are ignored. Failed factorizations propagate non-finite
    values on the active system.
    """
    right_hand_side = jnp.asarray(right_hand_side)
    if right_hand_side.shape != (self.capacity,):
        raise ValueError(
            "right-hand side must be a vector with the system capacity"
        )
    return self.factor.quadratic(right_hand_side)

logdet 

logdet() -> Scalar

Return the exact log determinant of the stabilized matrix.

Source code in src/gpex/linalg/dense.py 
131
132
133
def logdet(self) -> Scalar:
    """Return the exact log determinant of the stabilized matrix."""
    return self.factor.logdet()

covariance_reduction 

covariance_reduction(cross: Matrix) -> GramMatrix

Compute the covariance reduction, ignoring inactive cross rows.

Source code in src/gpex/linalg/dense.py 
135
136
137
def covariance_reduction(self, cross: Matrix) -> GramMatrix:
    """Compute the covariance reduction, ignoring inactive cross rows."""
    return self.project(cross).gram()

variance_reduction 

variance_reduction(cross: Matrix) -> ColumnValues

Reduce active cross rows without constructing a query-sized matrix.

Source code in src/gpex/linalg/dense.py 
139
140
141
def variance_reduction(self, cross: Matrix) -> ColumnValues:
    """Reduce active cross rows without constructing a query-sized matrix."""
    return self.project(cross).diagonal()

append 

append(update: LinearSystemUpdate) -> CholeskySystem

Activate reserved slots with a Schur-complement factor update.

Source code in src/gpex/linalg/dense.py 
143
144
145
146
147
148
def append(self, update: LinearSystemUpdate) -> CholeskySystem:
    """Activate reserved slots with a Schur-complement factor update."""
    cross, appended_covariance = self._trailing_blocks(update)
    with jax.default_matmul_precision("highest"):
        factor = self.factor.activate(update.indices, cross, appended_covariance)
    return eqx.tree_at(lambda s: s.factor, self, factor)

append_solve 

append_solve(
    update: LinearSystemUpdate,
    solution: Vector,
    right_hand_side: Vector,
) -> tuple[CholeskySystem, Vector]

Activate reserved slots and reuse the existing coefficients.

Inputs and outputs use canonical ordering. Old active RHS entries must be unchanged. Inactive inputs are ignored, and returned inactive coefficients are zero.

Source code in src/gpex/linalg/dense.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def append_solve(
    self,
    update: LinearSystemUpdate,
    solution: Vector,
    right_hand_side: Vector,
) -> tuple[CholeskySystem, Vector]:
    """Activate reserved slots and reuse the existing coefficients.

    Inputs and outputs use canonical ordering. Old active RHS entries
    must be unchanged. Inactive inputs are ignored, and returned inactive
    coefficients are zero.
    """
    cross, appended_covariance = self._trailing_blocks(update)
    with jax.default_matmul_precision("highest"):
        factor, solution = self.factor.activate_solve(
            update.indices,
            cross,
            appended_covariance,
            solution=solution,
            rhs=right_hand_side,
        )
    return eqx.tree_at(lambda s: s.factor, self, factor), solution

project 

project(cross: Matrix) -> CholeskyProjection

Cache a triangular projection, ignoring inactive cross rows.

Source code in src/gpex/linalg/dense.py 
173
174
175
def project(self, cross: Matrix) -> CholeskyProjection:
    """Cache a triangular projection, ignoring inactive cross rows."""
    return CholeskyProjection(self.factor.whiten(cross))

gpex.linalg.SolverConfiguration

Bases: ABC

Abstract backend preparing a numerical system with optional tiling.

prepare abstractmethod 

prepare(
    system: SymmetricLinearOperator,
    *,
    materialization_tile_size: int | None = None,
    active: Activity | None = None,
) -> LinearSystem
Source code in src/gpex/linalg/base.py 
47
48
49
50
51
52
53
54
@abstractmethod
def prepare(
    self,
    system: SymmetricLinearOperator,
    *,
    materialization_tile_size: int | None = None,
    active: Activity | None = None,
) -> LinearSystem: ...

gpex.linalg.LinearSystem

Bases: ABC

Abstract prepared system shared by inference and exact fitting.

Inputs and outputs use canonical observation order. Numerical state must remain differentiable and cross JIT boundaries as dynamic PyTree leaves. Invalid fitting trials retain non-finite quadratic forms and log determinants so optimizers can reject them without changing the system.

size abstractmethod property 

size: Count

Number of active scalar entries, as a dynamic JAX integer.

capacity abstractmethod property 

capacity: int

Static scalar dimension of input and output arrays, including padding.

solve abstractmethod 

solve(right_hand_side: RightHandSide) -> RightHandSide
Source code in src/gpex/linalg/base.py 
78
79
80
81
82
@abstractmethod
def solve(
    self,
    right_hand_side: RightHandSide,
) -> RightHandSide: ...

quadratic_form abstractmethod 

quadratic_form(right_hand_side: Vector) -> Scalar

Return rhs.T @ A⁻¹ @ rhs without forming the inverse.

Source code in src/gpex/linalg/base.py 
84
85
86
87
@abstractmethod
def quadratic_form(self, right_hand_side: Vector) -> Scalar:
    """Return ``rhs.T @ A⁻¹ @ rhs`` without forming the inverse."""
    ...

logdet abstractmethod 

logdet() -> Scalar

Return the log determinant of the prepared system.

Source code in src/gpex/linalg/base.py 
89
90
91
92
@abstractmethod
def logdet(self) -> Scalar:
    """Return the log determinant of the prepared system."""
    ...

covariance_reduction abstractmethod 

covariance_reduction(cross: Matrix) -> GramMatrix

Compute cross.T @ A⁻¹ @ cross.

Source code in src/gpex/linalg/base.py 
94
95
96
97
@abstractmethod
def covariance_reduction(self, cross: Matrix) -> GramMatrix:
    """Compute ``cross.T @ A⁻¹ @ cross``."""
    ...

variance_reduction abstractmethod 

variance_reduction(cross: Matrix) -> ColumnValues

Return only the diagonal of cross.T @ A⁻¹ @ cross.

Source code in src/gpex/linalg/base.py 
 99
100
101
102
@abstractmethod
def variance_reduction(self, cross: Matrix) -> ColumnValues:
    """Return only the diagonal of ``cross.T @ A⁻¹ @ cross``."""
    ...

project abstractmethod 

project(cross: Matrix) -> CovarianceProjection

Prepare reusable factors for consistent joint covariance reductions.

Source code in src/gpex/linalg/base.py 
104
105
106
107
@abstractmethod
def project(self, cross: Matrix) -> CovarianceProjection:
    """Prepare reusable factors for consistent joint covariance reductions."""
    ...

append abstractmethod 

append(update: LinearSystemUpdate) -> LinearSystem
Source code in src/gpex/linalg/base.py 
109
110
@abstractmethod
def append(self, update: LinearSystemUpdate) -> LinearSystem: ...

append_solve abstractmethod 

append_solve(
    update: LinearSystemUpdate,
    solution: Vector,
    right_hand_side: Vector,
) -> tuple[LinearSystem, Vector]
Source code in src/gpex/linalg/base.py 
112
113
114
115
116
117
118
@abstractmethod
def append_solve(
    self,
    update: LinearSystemUpdate,
    solution: Vector,
    right_hand_side: Vector,
) -> tuple[LinearSystem, Vector]: ...

gpex.linalg.SymmetricLinearOperator

Bases: ABC

Symmetric numerical input supporting matrix-free and dense operations.

Backends consume this interface without knowing about GP measurements. They choose when to materialize the matrix and how to prepare its state.

capacity abstractmethod property 

capacity: int

Materialized matrix dimension, including any reserved slots.

matmul abstractmethod 

matmul(
    right_hand_side: RightHandSide, *, tile_size: int = 128
) -> RightHandSide

Apply the system to canonical rows without materializing the matrix.

The input has shape (capacity, *columns). Reserved inactive slots are independent identity rows, so their input values pass through. tile_size bounds pairwise work, subject to indivisible operator outputs. Numerical inputs support JIT, vmap, and differentiation.

Source code in src/gpex/linalg/base.py 
134
135
136
137
138
139
140
141
142
143
144
145
@abstractmethod
def matmul(
    self, right_hand_side: RightHandSide, *, tile_size: int = 128
) -> RightHandSide:
    """Apply the system to canonical rows without materializing the matrix.

    The input has shape ``(capacity, *columns)``. Reserved inactive slots
    are independent identity rows, so their input values pass through.
    ``tile_size`` bounds pairwise work, subject to indivisible operator
    outputs. Numerical inputs support JIT, vmap, and differentiation.
    """
    ...

diagonal abstractmethod 

diagonal() -> Vector

Return canonical diagonal entries without a full matrix.

Source code in src/gpex/linalg/base.py 
147
148
149
150
@abstractmethod
def diagonal(self) -> Vector:
    """Return canonical diagonal entries without a full matrix."""
    ...

materialize abstractmethod 

materialize(*, tile_size: int | None) -> SquareMatrix

Construct the same system explicitly for dense methods.

Source code in src/gpex/linalg/base.py 
152
153
154
155
@abstractmethod
def materialize(self, *, tile_size: int | None) -> SquareMatrix:
    """Construct the same system explicitly for dense methods."""
    ...

gpex.linalg.LinearSystemUpdate

Bases: ABC

Activate unused slots without changing the system capacity.

indices are distinct canonical slots that were inactive. Cross rows use canonical ordering and are zero at previously inactive slots. Existing active entries of the system and right-hand side must remain unchanged.

indices instance-attribute 

indices: Indices

system instance-attribute 

system: SymmetricLinearOperator

previous_capacity abstractmethod property 

previous_capacity: int

Materialized dimension before the update, including reserved slots.

trailing_blocks abstractmethod 

trailing_blocks(
    *, tile_size: int | None
) -> tuple[CrossMatrix, AppendedMatrix]
Source code in src/gpex/linalg/base.py 
175
176
177
178
@abstractmethod
def trailing_blocks(
    self, *, tile_size: int | None
) -> tuple[CrossMatrix, AppendedMatrix]: ...

gpex.linalg.CovarianceProjection

Bases: ABC

Reusable reductions for targets of one prepared system.

Implementations own their numerical representation and expose it as dynamic PyTree leaves. Cross reductions require projections from the same system.

gram abstractmethod 

gram() -> GramMatrix

Return C.T @ A⁻¹ @ C for the projected cross-covariance C.

Source code in src/gpex/linalg/base.py 
188
189
190
191
@abstractmethod
def gram(self) -> GramMatrix:
    """Return ``C.T @ A⁻¹ @ C`` for the projected cross-covariance C."""
    ...

diagonal abstractmethod 

diagonal() -> ColumnValues

Return the reduction diagonal without a query-sized matrix.

Source code in src/gpex/linalg/base.py 
193
194
195
196
@abstractmethod
def diagonal(self) -> ColumnValues:
    """Return the reduction diagonal without a query-sized matrix."""
    ...

cross_gram abstractmethod 

cross_gram(other: Self) -> Matrix

Return C.T @ A⁻¹ @ D for two targets of the same system.

Source code in src/gpex/linalg/base.py 
198
199
200
201
@abstractmethod
def cross_gram(self, other: Self) -> Matrix:
    """Return ``C.T @ A⁻¹ @ D`` for two targets of the same system."""
    ...

gpex.linalg.CholeskyProjection

Bases: Module, CovarianceProjection

Factors for a symmetric bilinear covariance reduction.

Dense systems store factor = L⁻¹ C. Cross reductions require projections from the same system. All factors are dynamic leaves, including across an outer JIT boundary.

factor instance-attribute 

factor: Matrix

gram 

gram() -> GramMatrix

Return the full symmetric covariance reduction.

Source code in src/gpex/linalg/dense.py 
199
200
201
202
203
def gram(self) -> GramMatrix:
    """Return the full symmetric covariance reduction."""
    return jnp.matmul(
        self.factor.T, self.factor, precision=jax.lax.Precision.HIGHEST
    )

diagonal 

diagonal() -> ColumnValues

Return the reduction diagonal without a query-sized matrix.

Source code in src/gpex/linalg/dense.py 
205
206
207
208
209
def diagonal(self) -> ColumnValues:
    """Return the reduction diagonal without a query-sized matrix."""
    return jnp.einsum(
        "ij,ij->j", self.factor, self.factor, precision=jax.lax.Precision.HIGHEST
    )

cross_gram 

cross_gram(other: Self) -> Matrix

Return the cross reduction for two targets of the same system.

Source code in src/gpex/linalg/dense.py 
211
212
213
214
215
def cross_gram(self, other: Self) -> Matrix:
    """Return the cross reduction for two targets of the same system."""
    return jnp.matmul(
        self.factor.T, other.factor, precision=jax.lax.Precision.HIGHEST
    )

Moment assembly

gpex.moments assembles prior means and covariances of linear functionals in canonical block order. These functions accept public blocks or reusable PreparedFunctionals from gpex.moments.measurements. Preparation groups compatible operators and retains activity and ordering information.

assemble_covariance supports tiled evaluation while still returning the full dense matrix. assemble_covariance_diagonal evaluates only the diagonal. Inactive reserved entries are zero. These functions evaluate arrays and can be wrapped in an outer JIT with their consumers.

gpex.moments

Gaussian mean and covariance assembly for linear functionals.

apply_covariance 

apply_covariance(
    blocks: Functionals,
    kernel: ScalarKernel,
    right_hand_side: FunctionalColumns,
    *,
    tile_size: int = 128,
) -> FunctionalColumns

Apply latent covariance to canonical rows without constructing its matrix.

Accept vectors or arrays with leading dimension equal to the scalar functional size. Inactive inputs are ignored and inactive outputs are zero. One operator output is indivisible and can exceed tile_size. Pairwise tiles are recomputed in reverse mode, bounding working storage in both directions. Supports JIT, vmap, and differentiation of numerical inputs. General kernels still require quadratic pairwise work per application.

Source code in src/gpex/moments/assembly.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def apply_covariance(
    blocks: Functionals,
    kernel: ScalarKernel,
    right_hand_side: FunctionalColumns,
    *,
    tile_size: int = 128,
) -> FunctionalColumns:
    """Apply latent covariance to canonical rows without constructing its matrix.

    Accept vectors or arrays with leading dimension equal to the scalar
    functional size. Inactive inputs are ignored and inactive outputs are zero.
    One operator output is indivisible and can exceed ``tile_size``. Pairwise
    tiles are recomputed in reverse mode, bounding working storage in both
    directions. Supports JIT, vmap, and differentiation of numerical inputs.
    General kernels still require quadratic pairwise work per application.
    """
    _validate_tile_size(tile_size)
    measurements = prepare_functionals(blocks)
    _validate_kernel_compatibility((measurements,), kernel)
    right_hand_side = jnp.asarray(right_hand_side)
    if right_hand_side.ndim == 0 or right_hand_side.shape[0] != measurements.size:
        raise ValueError("right-hand side must start with the functional size")
    shape = right_hand_side.shape
    columns = measurements.to_grouped(measurements.mask(right_hand_side)).reshape(
        (measurements.size, prod(shape[1:]))
    )
    layouts = tuple(
        _TileLayout.from_group(group, tile_size) for group in measurements.groups
    )
    inputs = []
    start = 0
    for layout in layouts:
        inputs.append(
            jnp.pad(
                columns[start : start + layout.size],
                ((0, layout.padded_size - layout.size), (0, 0)),
            )
        )
        start += layout.size
    parts = []
    for left in layouts:
        product = jnp.zeros((left.size, columns.shape[1]))
        for right, values in zip(layouts, inputs, strict=True):
            product = product + _apply_group_covariance(left, right, kernel, values)
        parts.append(product)
    result = measurements.to_canonical(jnp.concatenate(parts))
    return measurements.mask(result).reshape(shape)

assemble_mean 

assemble_mean(
    blocks: Functionals, mean: ScalarFunction
) -> FlatMean

Apply each functional to the mean and return canonical scalar ordering.

Accept public blocks or prepared groups. Inactive reserved entries are zero. Supports JIT, vmap, and differentiation of numerical inputs.

Source code in src/gpex/moments/assembly.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def assemble_mean(blocks: Functionals, mean: ScalarFunction) -> FlatMean:
    """Apply each functional to the mean and return canonical scalar ordering.

    Accept public blocks or prepared groups. Inactive reserved entries are zero.
    Supports JIT, vmap, and differentiation of numerical inputs.
    """
    measurements = prepare_functionals(blocks)

    def scalar_mean(point: Point) -> Scalar:
        value = jnp.asarray(mean(point))
        if value.ndim != 0:
            raise ValueError("mean functions must return a scalar")
        return value

    parts = [
        eqx.filter_vmap(lambda operator, point: operator.apply(scalar_mean, point))(
            group.operators, group.points
        ).reshape(-1)
        for group in measurements.groups
    ]
    return measurements.mask(measurements.to_canonical(jnp.concatenate(parts)))

assemble_covariance 

assemble_covariance(
    left_blocks: Functionals,
    right_blocks: Functionals,
    kernel: ScalarKernel,
    *,
    tile_size: int | None = None,
) -> FlatCovariance

Materialize covariance, optionally bounding pairwise working tiles.

Tiles traverse numerical operator groups rather than public blocks. One operator output is never split. The final dense matrix is still allocated. With None, each complete group pair is vectorized.

Source code in src/gpex/moments/assembly.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def assemble_covariance(
    left_blocks: Functionals,
    right_blocks: Functionals,
    kernel: ScalarKernel,
    *,
    tile_size: int | None = None,
) -> FlatCovariance:
    """Materialize covariance, optionally bounding pairwise working tiles.

    Tiles traverse numerical operator groups rather than public blocks. One
    operator output is never split. The final dense matrix is still allocated.
    With ``None``, each complete group pair is vectorized.
    """
    if tile_size is not None:
        _validate_tile_size(tile_size)
    left, right = prepare_functionals(left_blocks), prepare_functionals(right_blocks)
    if left.input_dimension != right.input_dimension:
        raise ValueError("covariance blocks must use the same input dimension")
    _validate_kernel_compatibility((left, right), kernel)
    rows = []
    for group in left.groups:
        row = jnp.concatenate(
            [
                _assemble_group_covariance(group, other, kernel, tile_size=tile_size)
                for other in right.groups
            ],
            axis=1,
        )
        if len(left.groups) > 1 or len(right.groups) > 1:
            # TODO: Remove this temporary compiler/driver workaround once the
            # mixed CUDA regression passes without it on the fixed toolchain.
            # JAX/jaxlib 0.10.2 on H100 with NVIDIA 615.71.09 can raise an illegal
            # address when fusing values/slopes/averages covariance assembly.
            # Separate completed rows from the final join. Barriers on each
            # individual pair instead expose another failure at smaller sizes.
            row = jax.lax.optimization_barrier(row)
        rows.append(row)
    covariance = left.to_canonical(jnp.concatenate(rows, axis=0))
    return left.mask(right.mask(right.to_canonical(covariance.T)).T)

assemble_covariance_diagonal 

assemble_covariance_diagonal(
    blocks: Functionals, kernel: ScalarKernel
) -> FlatVector

Assemble the prior diagonal without allocating a full covariance.

Source code in src/gpex/moments/assembly.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def assemble_covariance_diagonal(
    blocks: Functionals, kernel: ScalarKernel
) -> FlatVector:
    """Assemble the prior diagonal without allocating a full covariance."""
    measurements = prepare_functionals(blocks)
    _validate_kernel_compatibility((measurements,), kernel)
    parts = []
    for group in measurements.groups:
        components = prod(group.output_shape)

        def point_diagonal(operator, point, components=components):
            covariance = operator_covariance(
                kernel, operator, operator, x=point, y=point
            )
            return jnp.diag(covariance.reshape((components, components)))

        parts.append(
            eqx.filter_vmap(point_diagonal)(group.operators, group.points).reshape(-1)
        )
    return measurements.mask(measurements.to_canonical(jnp.concatenate(parts)))

Covariance operator

CovarianceSystem describes an observation covariance including block noise, jitter, and reserved inactive slots. Backends can consume it through the SymmetricLinearOperator interface:

product = system.matmul(rhs, tile_size=128)
diagonal = system.diagonal()
matrix = system.materialize(tile_size=128)

rhs has shape (capacity, *columns) in canonical order. Multiplication accumulates tiles directly and passes inactive inputs through as identity rows. The diagonal contains ones in inactive slots. This differs from prepared solve, which ignores inactive inputs. Tile widths count scalar outputs, but one operator output is indivisible. Tile size is static under JIT.

Products support JIT, vmap, and numerical differentiation. Pairwise intermediates are recomputed in reverse mode so differentiation also avoids full covariance storage. General kernels still require quadratic pairwise work. Dense Cholesky remains the only backend, and likelihood evaluation remains exact.

gpex.covariance_systems.CovarianceSystem

Bases: Module, SymmetricLinearOperator

A stabilized GP covariance exposed as generic linear-algebra operations.

variance contains measurement-error variances in canonical scalar-output order.

measurements instance-attribute 

measurements: PreparedFunctionals

kernel instance-attribute 

kernel: ScalarKernel

variance instance-attribute 

variance: Vector

jitter class-attribute instance-attribute 

jitter: float = eqx.field(static=True)

capacity property 

capacity: int

diagonal_shift 

diagonal_shift() -> Vector

Return the positive diagonal observation shift.

Source code in src/gpex/covariance_systems.py 
47
48
49
def diagonal_shift(self) -> Vector:
    """Return the positive diagonal observation shift."""
    return jnp.where(self.measurements.active, self.variance + self.jitter, 1)

matmul 

matmul(
    right_hand_side: RightHandSide, *, tile_size: int = 128
) -> RightHandSide

Apply covariance, noise, and jitter using bounded pairwise tiles.

Inputs have shape (capacity, *columns) in canonical order. Inactive slots act as independent identity rows, unlike a prepared solve which ignores inactive inputs. Supports JIT, vmap, and differentiation.

Source code in src/gpex/covariance_systems.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def matmul(
    self, right_hand_side: RightHandSide, *, tile_size: int = 128
) -> RightHandSide:
    """Apply covariance, noise, and jitter using bounded pairwise tiles.

    Inputs have shape ``(capacity, *columns)`` in canonical order. Inactive
    slots act as independent identity rows, unlike a prepared solve which
    ignores inactive inputs. Supports JIT, vmap, and differentiation.
    """
    right_hand_side = jnp.asarray(right_hand_side)
    covariance_product = apply_covariance(
        self.measurements, self.kernel, right_hand_side, tile_size=tile_size
    )
    shift = self.diagonal_shift().reshape(
        (-1,) + (1,) * (right_hand_side.ndim - 1)
    )
    return covariance_product + shift * right_hand_side

diagonal 

diagonal() -> Vector

Return the stabilized diagonal, including unit inactive entries.

Source code in src/gpex/covariance_systems.py 
69
70
71
72
73
74
def diagonal(self) -> Vector:
    """Return the stabilized diagonal, including unit inactive entries."""
    return (
        assemble_covariance_diagonal(self.measurements, self.kernel)
        + self.diagonal_shift()
    )

materialize 

materialize(*, tile_size: int | None) -> SquareMatrix
Source code in src/gpex/covariance_systems.py 
76
77
78
79
80
81
82
83
84
def materialize(self, *, tile_size: int | None) -> SquareMatrix:
    covariance = assemble_covariance(
        self.measurements,
        self.measurements,
        self.kernel,
        tile_size=tile_size,
    )
    matrix = covariance + jnp.diag(self.diagonal_shift())
    return symmetrize(matrix)