Skip to content

Operators

See Linear operators for formulas and interpretations, and Implementing a custom operator for a complete extension example.

Custom operators

A LinearOperator implements output_shape(input_dimension) and apply(function, point). The point has shape (d,). For a scalar function, apply must return the declared output shape. It must also preserve any output axes already present on the supplied function, because covariance construction applies operators to both kernel arguments.

The operation must be linear in function. Its locations, weights, and other parameters must be independent of the function's values. Use JAX-compatible operations for jit, vmap, and differentiation. Store differentiable numerical parameters as array fields. Structural settings such as quadrature order can be static Equinox fields.

The default derivative_order is zero. Override it when the operator takes derivatives of the function. gpex checks it against the kernel's supported GP derivative order. A finite difference made from function evaluations has order zero under this contract.

Quadrature interfaces

QuadratureOperator.quadrature(point) returns absolute evaluation points with shape (q, d) and weights with shape (q,). Its inherited apply evaluates and weights the function while preserving existing output axes. The default scalar output_shape may be overridden for input-dimension restrictions.

Weights may be negative and need not sum to one. Include normalization and change-of-variable factors in the rule. Points and weights may depend on the query location and operator parameters. Keep the number of nodes fixed under tracing.

TensorQuadratureOperator.coordinate_quadrature(point) returns one (nodes, weights) pair per coordinate in input order. Each array is one-dimensional. Node counts can differ between coordinates and must remain fixed under tracing. The full weight is the product of the coordinate weights. Any override of quadrature or apply must represent that same product rule. This interface enables efficient covariance evaluation with SeparableKernel.

Built-in constructors and methods

Identity(), Gradient(), Hessian(), and Laplacian() are stateless. PartialDerivative(axis) selects one coordinate axis, which is static under JAX transformations. Numerical directions, displacements, radii, and smoothing scales are dynamic array parameters.

gpex.operators

Linear operators for observations and posterior queries.

LinearOperator

Bases: Module

Base class for linear measurements and query operators.

Subclasses declare the result produced at one input location with output_shape and implement the linear action with apply. apply must remain linear in function and use JAX-compatible operations.

derivative_order property 

derivative_order: int

Required derivative order of the latent function.

The default, zero, is for operators using only function values. Custom differential operators must override this so covariance compilation can check the kernel's smoothness before applying automatic differentiation.

output_shape abstractmethod 

output_shape(input_dimension: int) -> tuple[int, ...]

Return the operator output shape at one input location.

Source code in src/gpex/operators/base.py 
38
39
40
@abstractmethod
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    """Return the operator output shape at one input location."""

apply abstractmethod 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue

Apply the operator to function at point.

Source code in src/gpex/operators/base.py 
42
43
44
@abstractmethod
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    """Apply the operator to ``function`` at ``point``."""

QuadratureOperator

Bases: LinearOperator

A scalar linear measurement defined by weighted function evaluations.

Subclasses implement quadrature(point) to return evaluation points of shape (q, d) and weights of shape (q,). Both may depend on the query location and operator parameters, but not on the function being evaluated. Weights may be signed and are not normalized by this class.

The rule must use JAX-compatible operations and keep q fixed under tracing. apply supports JIT, batching, and differentiation through the points and weights. It also preserves trailing function-output axes when operators are nested during covariance compilation.

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]

Return scalar output; subclasses may validate the input dimension.

Source code in src/gpex/operators/quadrature.py 
42
43
44
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    """Return scalar output; subclasses may validate the input dimension."""
    return ()

quadrature abstractmethod 

quadrature(point: Point) -> QuadratureRule

Return absolute evaluation points (q, d) and weights (q,).

Include the rule's normalization and any change-of-variable factors in the weights. The point count and input dimension are static.

Source code in src/gpex/operators/quadrature.py 
46
47
48
49
50
51
52
@abstractmethod
def quadrature(self, point: Point) -> QuadratureRule:
    """Return absolute evaluation points ``(q, d)`` and weights ``(q,)``.

    Include the rule's normalization and any change-of-variable factors
    in the weights. The point count and input dimension are static.
    """

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/quadrature.py 
54
55
56
57
58
59
60
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    self.output_shape(point.shape[-1])
    points, weights = self.quadrature(point)
    values = jax.vmap(function)(points)
    return jnp.tensordot(
        weights, values, axes=(0, 0), precision=jax.lax.Precision.HIGHEST
    )

TensorQuadratureOperator

Bases: QuadratureOperator

A scalar quadrature rule that factors over input coordinates.

Implement coordinate_quadrature(point) to return one pair of node and weight vectors per coordinate. Nodes are absolute coordinate values, and full weights are products of coordinate weights. Both may depend on the entire query point and operator parameters, independently of the function. Node counts are static under tracing and may differ between coordinates.

The inherited apply evaluates the Cartesian-product rule. Covariance evaluation with a SeparableKernel contracts coordinate rules directly, preserving the same approximation without constructing the full grid. Overrides of apply or quadrature must represent that same rule.

coordinate_quadrature abstractmethod 

coordinate_quadrature(point: Point) -> TensorQuadratureRule

Return ((nodes_0, weights_0), ...) in input-coordinate order.

Each pair contains equally sized one-dimensional arrays. Include normalization and change-of-variable factors in the weights; signed and unnormalized rules are supported.

Source code in src/gpex/operators/quadrature.py 
78
79
80
81
82
83
84
85
@abstractmethod
def coordinate_quadrature(self, point: Point) -> TensorQuadratureRule:
    """Return ``((nodes_0, weights_0), ...)`` in input-coordinate order.

    Each pair contains equally sized one-dimensional arrays. Include
    normalization and change-of-variable factors in the weights; signed
    and unnormalized rules are supported.
    """

quadrature 

quadrature(point: Point) -> QuadratureRule
Source code in src/gpex/operators/quadrature.py 
87
88
89
90
91
92
93
94
95
96
97
98
99
def quadrature(self, point: Point) -> QuadratureRule:
    rules = self.coordinate_quadrature(point)
    points = jnp.stack(
        jnp.meshgrid(*(nodes for nodes, _ in rules), indexing="ij"), axis=-1
    ).reshape(-1, point.shape[-1])
    weights = jnp.prod(
        jnp.stack(
            jnp.meshgrid(*(weights for _, weights in rules), indexing="ij"),
            axis=-1,
        ),
        axis=-1,
    ).reshape(-1)
    return points, weights

Identity

Bases: LinearOperator

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
16
17
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    return ()

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
19
20
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    return function(point)

Increment 

Increment(displacement: Displacement)

Bases: LinearOperator

Return the unnormalized increment f(x + displacement) - f(x).

Source code in src/gpex/operators/differences.py 
18
19
20
21
22
def __init__(self, displacement: Displacement) -> None:
    displacement = jnp.asarray(displacement)
    if displacement.ndim != 1 or displacement.size == 0:
        raise ValueError("displacement must be a non-empty one-dimensional array")
    self.displacement = displacement

displacement instance-attribute 

displacement: Displacement = displacement

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/differences.py 
24
25
26
27
28
29
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if self.displacement.shape != (input_dimension,):
        raise ValueError(
            "displacement shape must match the locations' input dimension"
        )
    return ()

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/differences.py 
31
32
33
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    self.output_shape(point.shape[-1])
    return function(point + self.displacement) - function(point)

PartialDerivative

Bases: LinearOperator

axis class-attribute instance-attribute 

axis: int = eqx.field(static=True)

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
30
31
32
33
34
35
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if not 0 <= self.axis < input_dimension:
        raise ValueError(
            f"axis {self.axis} is invalid for dimension {input_dimension}"
        )
    return ()

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
37
38
39
40
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    self.output_shape(point.shape[-1])
    tangent = jnp.zeros_like(point).at[self.axis].set(1)
    return jax.jvp(function, (point,), (tangent,))[1]

DirectionalDerivative 

DirectionalDerivative(direction: Direction)

Bases: LinearOperator

Source code in src/gpex/operators/local.py 
100
101
102
103
104
def __init__(self, direction: Direction) -> None:
    direction = jnp.asarray(direction)
    if direction.ndim != 1:
        raise ValueError("direction must be a one-dimensional array")
    self.direction = direction

direction instance-attribute 

direction: Direction = direction

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
110
111
112
113
114
115
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if self.direction.shape != (input_dimension,):
        raise ValueError(
            "direction shape must match the locations' input dimension"
        )
    return ()

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
117
118
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    return jax.jvp(function, (point,), (self.direction,))[1]

Gradient

Bases: LinearOperator

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
48
49
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    return (input_dimension,)

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
51
52
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    return jax.jacfwd(function)(point)

Hessian

Bases: LinearOperator

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
60
61
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    return (input_dimension, input_dimension)

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
63
64
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    return jax.jacfwd(jax.jacfwd(function))(point)

Laplacian

Bases: LinearOperator

Return the trace of the input Hessian at each location.

Sum second coordinate derivatives without constructing the full Hessian.

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/local.py 
77
78
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    return ()

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/local.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    def second_derivative(axis: int) -> FunctionalValue:
        tangent = jax.nn.one_hot(axis, point.shape[0])

        def first_derivative(variable: Point) -> FunctionalValue:
            return jax.jvp(function, (variable,), (tangent,))[1]

        return jax.jvp(first_derivative, (point,), (tangent,))[1]

    def accumulate(axis: int, total: FunctionalValue) -> FunctionalValue:
        return total + second_derivative(axis)

    # Preserve array-valued intermediate functions from another operator.
    initial = jnp.zeros_like(function(point))
    return jax.lax.fori_loop(0, point.shape[0], accumulate, initial)

BoxAverage 

BoxAverage(
    radius: BoxRadius | float, *, quadrature_order: int = 8
)

Bases: TensorQuadratureOperator

Average a scalar function over a box centered at each input location.

A scalar radius defines a square or hypercube. A radius vector defines one half-width per input dimension. The integral is approximated by a symmetric tensor-product midpoint rule with quadrature_order ** dimension points.

Construction and evaluation support JIT, batching, and differentiation of finite positive radius values. Quadrature order remains static.

Source code in src/gpex/operators/quadrature.py 
116
117
118
119
120
121
122
123
124
125
126
127
def __init__(self, radius: BoxRadius | float, *, quadrature_order: int = 8) -> None:
    radius = jnp.asarray(radius)
    if radius.ndim > 1 or radius.size == 0:
        raise ValueError("box-average radius must be scalar or one-dimensional")
    if quadrature_order < 1:
        raise ValueError("quadrature order must be positive")
    self.radius = error_if(
        radius,
        jnp.any(~jnp.isfinite(radius) | (radius <= 0)),
        "box-average radius must be finite and strictly positive",
    )
    self.quadrature_order = quadrature_order

radius instance-attribute 

radius: BoxRadius = error_if(
    radius,
    jnp.any(~jnp.isfinite(radius) | (radius <= 0)),
    "box-average radius must be finite and strictly positive",
)

quadrature_order class-attribute instance-attribute 

quadrature_order: int = quadrature_order

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/quadrature.py 
129
130
131
132
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if self.radius.ndim == 1 and self.radius.shape != (input_dimension,):
        raise ValueError("radius shape must match the locations' input dimension")
    return ()

coordinate_quadrature 

coordinate_quadrature(point: Point) -> TensorQuadratureRule
Source code in src/gpex/operators/quadrature.py 
134
135
136
137
138
139
140
141
142
143
144
145
def coordinate_quadrature(self, point: Point) -> TensorQuadratureRule:
    input_dimension = point.shape[-1]
    radius = jnp.broadcast_to(self.radius, (input_dimension,))
    normalized_nodes = (
        2.0 * (jnp.arange(self.quadrature_order) + 0.5) / self.quadrature_order
        - 1.0
    )
    weights = jnp.full((self.quadrature_order,), 1.0 / self.quadrature_order)
    return tuple(
        (point[axis] + normalized_nodes * radius[axis], weights)
        for axis in range(input_dimension)
    )

GaussianSmoothing 

GaussianSmoothing(
    scale: GaussianScale | float,
    *,
    quadrature_order: int = 8,
)

Bases: TensorQuadratureOperator

Average a scalar function under a centered Gaussian density.

A scalar scale defines an isotropic Gaussian. A scale vector defines one standard deviation per input dimension. The expectation is approximated with tensor-product Gauss--Hermite quadrature using quadrature_order ** dimension points.

Construction and evaluation support JIT, batching, and differentiation of finite positive scale values. Quadrature order remains static.

Source code in src/gpex/operators/quadrature.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def __init__(
    self, scale: GaussianScale | float, *, quadrature_order: int = 8
) -> None:
    scale = jnp.asarray(scale)
    if scale.ndim > 1 or scale.size == 0:
        raise ValueError(
            "Gaussian-smoothing scale must be scalar or one-dimensional"
        )
    nodes, weights = np.polynomial.hermite.hermgauss(quadrature_order)
    self.scale = error_if(
        scale,
        jnp.any(~jnp.isfinite(scale) | (scale <= 0)),
        "Gaussian-smoothing scale must be finite and strictly positive",
    )
    self.quadrature_order = quadrature_order
    self._nodes = tuple(float(node * np.sqrt(2.0)) for node in nodes)
    self._weights = tuple(float(weight / np.sqrt(np.pi)) for weight in weights)

scale instance-attribute 

scale: GaussianScale = error_if(
    scale,
    jnp.any(~jnp.isfinite(scale) | (scale <= 0)),
    "Gaussian-smoothing scale must be finite and strictly positive",
)

quadrature_order class-attribute instance-attribute 

quadrature_order: int = quadrature_order

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/quadrature.py 
183
184
185
186
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if self.scale.ndim == 1 and self.scale.shape != (input_dimension,):
        raise ValueError("scale shape must match the locations' input dimension")
    return ()

coordinate_quadrature 

coordinate_quadrature(point: Point) -> TensorQuadratureRule
Source code in src/gpex/operators/quadrature.py 
188
189
190
191
192
193
194
195
196
197
def coordinate_quadrature(self, point: Point) -> TensorQuadratureRule:
    input_dimension = point.shape[-1]
    scale = jnp.broadcast_to(self.scale, (input_dimension,))
    nodes = jnp.asarray(self._nodes)
    weights = jnp.asarray(self._weights)

    return tuple(
        (point[axis] + nodes * scale[axis], weights)
        for axis in range(input_dimension)
    )

CumulativeIntegral1D 

CumulativeIntegral1D(
    base: IntegralBase | float, *, quadrature_order: int = 8
)

Bases: QuadratureOperator

Integrate a one-dimensional function from base to each location.

Gauss--Legendre quadrature is mapped from [-1, 1] to the oriented interval from base to x. Locations below base therefore produce the correctly signed integral, and the value at base is exactly zero.

Source code in src/gpex/operators/quadrature.py 
213
214
215
216
217
218
219
220
221
222
223
def __init__(
    self, base: IntegralBase | float, *, quadrature_order: int = 8
) -> None:
    base = jnp.asarray(base)
    if base.ndim != 0:
        raise ValueError("integration base must be scalar")
    nodes, weights = np.polynomial.legendre.leggauss(quadrature_order)
    self.base = base
    self.quadrature_order = quadrature_order
    self._nodes = tuple(float(node) for node in nodes)
    self._weights = tuple(float(weight) for weight in weights)

base instance-attribute 

base: IntegralBase = base

quadrature_order class-attribute instance-attribute 

quadrature_order: int = quadrature_order

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/quadrature.py 
225
226
227
228
229
230
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if input_dimension != 1:
        raise ValueError(
            "CumulativeIntegral1D only supports one-dimensional inputs"
        )
    return ()

quadrature 

quadrature(point: Point) -> QuadratureRule
Source code in src/gpex/operators/quadrature.py 
232
233
234
235
236
237
238
def quadrature(self, point: Point) -> QuadratureRule:
    nodes = jnp.asarray(self._nodes)
    weights = jnp.asarray(self._weights)
    half_interval = (point[0] - self.base) / 2.0
    midpoint = (point[0] + self.base) / 2.0
    quadrature_points = midpoint + half_interval * nodes
    return quadrature_points[:, None], half_interval * weights

HaarWavelet1D 

HaarWavelet1D(
    scale: WaveletScale | float,
    *,
    quadrature_order: int = 8,
)

Bases: QuadratureOperator

A translated, fixed-scale Haar-wavelet coefficient in one dimension.

scale is the full width of the wavelet support. The analyzing wavelet is positive on the left half, negative on the right half, has unit \(L^2\) norm, and therefore annihilates constants. Each half-integral is approximated with a symmetric midpoint rule.

Construction and evaluation support JIT, batching, and differentiation of finite positive scale values. Quadrature order remains static.

Source code in src/gpex/operators/quadrature.py 
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def __init__(
    self, scale: WaveletScale | float, *, quadrature_order: int = 8
) -> None:
    scale = jnp.asarray(scale)
    if scale.ndim != 0:
        raise ValueError("Haar-wavelet scale must be scalar")
    if quadrature_order < 1:
        raise ValueError("quadrature order must be positive")
    self.scale = error_if(
        scale,
        ~jnp.isfinite(scale) | (scale <= 0),
        "Haar-wavelet scale must be finite and strictly positive",
    )
    self.quadrature_order = quadrature_order

scale instance-attribute 

scale: WaveletScale = error_if(
    scale,
    ~jnp.isfinite(scale) | (scale <= 0),
    "Haar-wavelet scale must be finite and strictly positive",
)

quadrature_order class-attribute instance-attribute 

quadrature_order: int = quadrature_order

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/quadrature.py 
271
272
273
274
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    if input_dimension != 1:
        raise ValueError("HaarWavelet1D currently supports one-dimensional inputs")
    return ()

quadrature 

quadrature(point: Point) -> QuadratureRule
Source code in src/gpex/operators/quadrature.py 
276
277
278
279
280
281
282
283
284
285
286
def quadrature(self, point: Point) -> QuadratureRule:
    nodes = (jnp.arange(self.quadrature_order) + 0.5) / self.quadrature_order
    half_width = self.scale / 2.0
    left_offsets = -half_width + nodes * half_width
    right_offsets = nodes * half_width
    offsets = jnp.concatenate((left_offsets, right_offsets))[:, None]
    signs = jnp.concatenate(
        (jnp.ones(self.quadrature_order), -jnp.ones(self.quadrature_order))
    )
    weights = signs * jnp.sqrt(self.scale) / (2 * self.quadrature_order)
    return point + offsets, weights

LinearCombination 

LinearCombination(
    operators: tuple[LinearOperator, ...]
    | list[LinearOperator],
    coefficients: OperatorCoefficients,
)

Bases: LinearOperator

A weighted sum of operators with a common output shape.

Source code in src/gpex/operators/base.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    operators: tuple[LinearOperator, ...] | list[LinearOperator],
    coefficients: OperatorCoefficients,
) -> None:
    operators = tuple(operators)
    coefficients = jnp.asarray(coefficients)
    if not operators:
        raise ValueError("at least one operator is required")
    if coefficients.shape != (len(operators),):
        raise ValueError("coefficients must contain one value per operator")
    self.operators = operators
    self.coefficients = coefficients

operators instance-attribute 

operators: tuple[LinearOperator, ...] = operators

coefficients instance-attribute 

coefficients: OperatorCoefficients = coefficients

derivative_order property 

derivative_order: int

output_shape 

output_shape(input_dimension: int) -> tuple[int, ...]
Source code in src/gpex/operators/base.py 
71
72
73
74
75
def output_shape(self, input_dimension: int) -> tuple[int, ...]:
    shapes = [operator.output_shape(input_dimension) for operator in self.operators]
    if any(shape != shapes[0] for shape in shapes[1:]):
        raise ValueError("combined operators must have the same output shape")
    return shapes[0]

apply 

apply(
    function: PointFunction, point: Point
) -> FunctionalValue
Source code in src/gpex/operators/base.py 
77
78
79
80
81
82
def apply(self, function: PointFunction, point: Point) -> FunctionalValue:
    values = jnp.stack(
        [operator.apply(function, point) for operator in self.operators]
    )
    coefficient_shape = (len(self.operators),) + (1,) * (values.ndim - 1)
    return jnp.sum(values * self.coefficients.reshape(coefficient_shape), axis=0)