Skip to content

Blocks

gpex.blocks

Structured functional, observation-design, and data blocks.

Locations 

Locations(points: Points)

Bases: Module

Non-empty real locations in JAX's default floating precision.

Parameters:

Name Type Description Default
points Points

Floating array of shape (n, d): n locations, each with d coordinates. Callers supply floating values. For a one-dimensional vector x of shape (n,), use x[:, None] to add the input axis.

 required
Source code in src/gpex/blocks/functionals.py 
30
31
32
33
34
35
36
def __init__(self, points: Points) -> None:
    points = jnp.asarray(points)
    if points.ndim != 2:
        raise ValueError("locations must have shape (n, d)")
    if points.shape[0] == 0 or points.shape[1] == 0:
        raise ValueError("locations cannot have empty axes")
    self.points = points

points instance-attribute 

points: Points = points

count property 

count: int

input_dimension property 

input_dimension: int

FunctionalBlock 

FunctionalBlock(
    locations: Locations | Points,
    operator: LinearOperator,
    name: str,
)

Bases: Module

Named linear outputs at a set of locations, without observed values.

Parameters:

Name Type Description Default
locations Locations | Points

Floating coordinates with shape (n, d), or Locations.

 required
operator LinearOperator

Linear operator. Its output shape determines each location's outputs: () for values, (d,) for gradients, for example.

 required
name str

Nonempty block name used to access results.

 required
Source code in src/gpex/blocks/functionals.py 
62
63
64
65
66
67
68
69
70
71
72
73
def __init__(
    self, locations: Locations | Points, operator: LinearOperator, name: str
) -> None:
    locations = (
        locations if isinstance(locations, Locations) else Locations(locations)
    )
    if not name:
        raise ValueError("block names cannot be empty")
    operator.output_shape(locations.input_dimension)
    self.locations = locations
    self.operator = operator
    self.name = name

locations instance-attribute 

locations: Locations = locations

operator instance-attribute 

operator: LinearOperator = operator

name class-attribute instance-attribute 

name: str = name

shape property 

shape: tuple[int, ...]

size property 

size: int

append 

append(locations: Locations) -> FunctionalBlock

Return this block extended by locations, retaining its operator and name.

The new locations must have the same input dimension. Supports JIT and differentiation; the original block is unchanged.

Source code in src/gpex/blocks/functionals.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
def append(self, locations: Locations) -> FunctionalBlock:
    """Return this block extended by locations, retaining its operator and name.

    The new locations must have the same input dimension. Supports JIT and
    differentiation; the original block is unchanged.
    """
    if locations.input_dimension != self.locations.input_dimension:
        raise ValueError("appended locations must have the same input dimension")
    return FunctionalBlock(
        Locations(jnp.concatenate((self.locations.points, locations.points))),
        self.operator,
        self.name,
    )

FunctionalQuery 

FunctionalQuery(blocks: Sequence[FunctionalBlock])

Bases: Module

An ordered collection of named outputs to predict or sample.

Parameters:

Name Type Description Default
blocks Sequence[FunctionalBlock]

Nonempty sequence of FunctionalBlock objects with unique names and a common input dimension. Ordering defines the flat output.

 required
Source code in src/gpex/blocks/queries.py 
24
25
26
27
28
def __init__(self, blocks: Sequence[FunctionalBlock]) -> None:
    blocks = tuple(blocks)
    _validate_collection(blocks)
    self.blocks = blocks
    self.index_map = IndexMap.from_blocks(blocks)

blocks instance-attribute 

blocks: tuple[FunctionalBlock, ...] = blocks

index_map class-attribute instance-attribute 

index_map: IndexMap = IndexMap.from_blocks(blocks)

functional_blocks property 

functional_blocks: tuple[FunctionalBlock, ...]

ObservationDesignBlock 

ObservationDesignBlock(
    functional: FunctionalBlock,
    variance: PositiveParameter | Scalar | float,
    *,
    max_size: int | None = None,
    count: Count | int | None = None,
)

Bases: Module

A functional block with measurement noise, but no observed outcome.

Parameters:

Name Type Description Default
functional FunctionalBlock

Named outputs to measure, including their locations and operator.

 required
variance PositiveParameter | Scalar | float

One finite, nonnegative scalar measurement-error variance, shared by every location and output component. A scalar array or float is fixed. A scalar PositiveParameter is trainable and preserves positivity during fitting. Observation errors are independent.

 required
max_size int | None

Reserved number of locations, at least the supplied number. Defaults to max(1024, supplied locations). Set it explicitly to the supplied count for compact storage. Operator output components each occupy a scalar slot per reserved location.

 None
count Count | int | None

Active prefix length. Defaults to all supplied locations. May be traced, unlike capacity. Inactive locations repeat the first supplied point, which must be valid for the operator and kernel even when the active count is zero.

 None
Source code in src/gpex/blocks/designs.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def __init__(
    self,
    functional: FunctionalBlock,
    variance: PositiveParameter | Scalar | float,
    *,
    max_size: int | None = None,
    count: Count | int | None = None,
) -> None:
    if isinstance(variance, PositiveParameter):
        value = variance.value
    else:
        variance = jnp.asarray(variance)
        variance = error_if(
            variance,
            jnp.any(~jnp.isfinite(variance) | (variance < 0)),
            "noise variances must be finite and nonnegative",
        )
        value = variance
    if value.ndim != 0:
        raise ValueError("noise variance must be scalar")
    supplied = functional.locations.count
    capacity = max(1024, supplied) if max_size is None else max_size
    if isinstance(capacity, bool) or not isinstance(capacity, int):
        raise TypeError("max_size must be an integer")
    if capacity < supplied:
        raise ValueError("max_size must be at least the supplied location count")
    count = jnp.asarray(supplied if count is None else count)
    if count.ndim != 0 or not jnp.issubdtype(count.dtype, jnp.integer):
        raise ValueError("active count must be an integer scalar")
    count = error_if(
        count,
        (count < 0) | (count > supplied),
        "active count must lie within the supplied locations",
    )
    points = jnp.pad(
        functional.locations.points,
        ((0, capacity - supplied), (0, 0)),
        mode="edge",
    )
    points = jnp.where((jnp.arange(capacity) < count)[:, None], points, points[0])
    self.functional = eqx.tree_at(lambda f: f.locations.points, functional, points)
    self.variance = variance
    self.count = count

functional instance-attribute 

functional: FunctionalBlock = eqx.tree_at(
    lambda f: f.locations.points, functional, points
)

variance instance-attribute 

variance: PositiveParameter | Scalar = variance

count instance-attribute 

count: Count = count

capacity property 

capacity: int

Reserved locations, independent of the active count.

active property 

active: Activity

active_flat property 

active_flat: FlatActivity

Activity in canonical scalar-output order.

active_size property 

active_size: Count

Number of active scalar observations, including operator components.

variances property 

variances: VarianceValues

The shared variance repeated over the functional output shape.

append 

append(locations: Locations) -> ObservationDesignBlock

Fill unused locations, retaining operator, name, variance and capacity.

Locations have shape (m, d) with the existing input dimension d. The returned design supports JIT and differentiation.

Source code in src/gpex/blocks/designs.py 
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
def append(self, locations: Locations) -> ObservationDesignBlock:
    """Fill unused locations, retaining operator, name, variance and capacity.

    Locations have shape ``(m, d)`` with the existing input dimension ``d``.
    The returned design supports JIT and differentiation.
    """
    if locations.input_dimension != self.functional.locations.input_dimension:
        raise ValueError("appended locations must have the same input dimension")
    if locations.count > self.capacity:
        raise ValueError("appended locations exceed max_size")
    points, count = error_if(
        (self.functional.locations.points, self.count),
        self.count + locations.count > self.capacity,
        "appended locations exceed max_size",
    )
    points = jax.lax.dynamic_update_slice_in_dim(
        points, locations.points, count, axis=0
    )
    count = count + locations.count
    points = jnp.where(
        (jnp.arange(self.capacity) < count)[:, None], points, points[0]
    )
    return eqx.tree_at(
        lambda b: (b.functional.locations.points, b.count), self, (points, count)
    )

replace 

replace(points: Points) -> ObservationDesignBlock

Replace all active locations, preserving capacity and shared variance.

Empty contents retain the previous valid padding point. Supports JIT and differentiation. Different input shapes specialize this operation, but its result always has the original reserved shapes.

Source code in src/gpex/blocks/designs.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def replace(self, points: Points) -> ObservationDesignBlock:
    """Replace all active locations, preserving capacity and shared variance.

    Empty contents retain the previous valid padding point. Supports JIT
    and differentiation. Different input shapes specialize this operation,
    but its result always has the original reserved shapes.
    """
    points = jnp.asarray(points)
    expected = self.functional.locations.input_dimension
    if points.ndim != 2 or points.shape[1] != expected:
        raise ValueError(f"replacement locations must have shape (n, {expected})")
    count = points.shape[0]
    if count > self.capacity:
        raise ValueError("replacement locations exceed max_size")
    if count == 0:
        points = self.functional.locations.points[:1]
    functional = FunctionalBlock(
        points, self.functional.operator, self.functional.name
    )
    return ObservationDesignBlock(
        functional, self.variance, max_size=self.capacity, count=count
    )

ObservationDesign 

ObservationDesign(blocks: Sequence[ObservationDesignBlock])

Bases: Module

An ordered measurement design without observed outcomes.

Parameters:

Name Type Description Default
blocks Sequence[ObservationDesignBlock]

Nonempty sequence of design blocks with unique names and a common input dimension. Variances are flattened in the same order as outputs.

 required
Source code in src/gpex/blocks/designs.py 
184
185
186
187
188
def __init__(self, blocks: Sequence[ObservationDesignBlock]) -> None:
    blocks = tuple(blocks)
    self.blocks = blocks
    _validate_collection(self.functional_blocks)
    self.index_map = IndexMap.from_blocks(self.functional_blocks)

blocks instance-attribute 

blocks: tuple[ObservationDesignBlock, ...] = blocks

index_map class-attribute instance-attribute 

index_map: IndexMap = IndexMap.from_blocks(
    self.functional_blocks
)

query property 

query: FunctionalQuery

functional_blocks property 

functional_blocks: tuple[FunctionalBlock, ...]

variance_flat property 

variance_flat: ObservationVector

Measurement-error variances in canonical scalar-output order.

active_size property 

active_size: Count

ObservationBlock 

ObservationBlock(
    design: ObservationDesignBlock, y: FunctionalValues
)

Bases: Module

A measurement design and its observed outputs.

Parameters:

Name Type Description Default
design ObservationDesignBlock

Functional outputs and their independent Gaussian noise variances.

 required
y FunctionalValues

Observed values with shape (active_count, *output_shape): (n,) for values or (n, d) for gradients, for example. A full capacity buffer is also accepted. Inactive values are zeroed.

 required
Source code in src/gpex/blocks/datasets.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(self, design: ObservationDesignBlock, y: FunctionalValues) -> None:
    y = jnp.asarray(y)
    functional = design.functional
    shape = functional.shape
    if y.ndim != len(shape) or y.shape[1:] != shape[1:] or y.shape[0] > shape[0]:
        raise ValueError(
            f"values for {functional.name!r} have shape {y.shape}; expected {shape}"
        )
    if y.shape[0] < design.capacity:
        y = error_if(
            y,
            design.count != y.shape[0],
            "values must match the active location count",
        )
    y = jnp.pad(y, ((0, design.capacity - y.shape[0]), *((0, 0),) * (y.ndim - 1)))
    self.design = design
    self.y = jnp.where(design.active.reshape((-1,) + (1,) * (y.ndim - 1)), y, 0)

design instance-attribute 

design: ObservationDesignBlock = design

y instance-attribute 

y: FunctionalValues = jnp.where(
    design.active.reshape((-1,) + (1,) * (y.ndim - 1)),
    y,
    0,
)

functional property 

functional: FunctionalBlock

The stored functional description, without noise or outcomes.

capacity property 

capacity: int

count property 

count: Count

from_arrays classmethod 

from_arrays(
    points: Points,
    y: FunctionalValues,
    *,
    variance: PositiveParameter | Scalar | float,
    name: str,
    operator: LinearOperator | None = None,
    max_size: int | None = None,
    padding_point: Point | None = None,
) -> ObservationBlock

Construct observations directly from coordinates and outcomes.

points has shape (n, d). The default identity operator expects y of shape (n,); otherwise use (n, *operator.output_shape(d)). variance is one scalar measurement-error variance shared by the block. Numerical values may be traced; the name and shapes must remain static.

max_size reserves locations, defaulting to max(1024, n). Unused slots repeat the first point. For n=0, supply a padding_point valid for the kernel, mean and operator, and an array points with shape (0, d). Padding is excluded from inference.

Source code in src/gpex/blocks/datasets.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
@classmethod
def from_arrays(
    cls,
    points: Points,
    y: FunctionalValues,
    *,
    variance: PositiveParameter | Scalar | float,
    name: str,
    operator: LinearOperator | None = None,
    max_size: int | None = None,
    padding_point: Point | None = None,
) -> ObservationBlock:
    """Construct observations directly from coordinates and outcomes.

    ``points`` has shape ``(n, d)``. The default identity operator expects
    ``y`` of shape ``(n,)``; otherwise use ``(n, *operator.output_shape(d))``.
    ``variance`` is one scalar measurement-error variance shared by the block.
    Numerical values may be traced; the name and shapes must remain static.

    ``max_size`` reserves locations, defaulting to ``max(1024, n)``.
    Unused slots repeat the first point. For ``n=0``, supply a
    ``padding_point`` valid for the kernel, mean and operator, and an array
    ``points`` with shape ``(0, d)``. Padding is excluded from inference.
    """
    operator = Identity() if operator is None else operator
    points = jnp.asarray(points)
    if points.ndim != 2:
        raise ValueError("locations must have shape (n, d)")
    count = points.shape[0]
    if count == 0:
        if padding_point is None:
            raise ValueError("empty observations require a valid padding_point")
        padding_point = jnp.asarray(padding_point)
        if padding_point.shape != (points.shape[1],):
            raise ValueError("padding_point must match the input dimension")
        points = padding_point[None, :]
    return cls(
        ObservationDesignBlock(
            FunctionalBlock(points, operator, name),
            variance,
            max_size=max_size,
            count=count,
        ),
        y,
    )

replace 

replace(
    points: Points, y: FunctionalValues
) -> ObservationBlock

Replace contents within capacity, retaining operator, name and variance.

Source code in src/gpex/blocks/datasets.py 
118
119
120
def replace(self, points: Points, y: FunctionalValues) -> ObservationBlock:
    """Replace contents within capacity, retaining operator, name and variance."""
    return ObservationBlock(self.design.replace(points), y)

append 

append(
    locations: Locations, y: FunctionalValues
) -> ObservationBlock

Append outcomes using this block's operator and shared variance.

Locations have shape (m, d) with the existing input dimension d; y must have shape (m, *output_shape). Returns a new block and supports JIT and differentiation, including the shared variance parameter.

Source code in src/gpex/blocks/datasets.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def append(self, locations: Locations, y: FunctionalValues) -> ObservationBlock:
    """Append outcomes using this block's operator and shared variance.

    Locations have shape ``(m, d)`` with the existing input dimension ``d``;
    ``y`` must have shape ``(m, *output_shape)``. Returns a new block and
    supports JIT and differentiation, including the shared variance parameter.
    """
    design = self.design.append(locations)
    y = jnp.asarray(y)
    expected_shape = (locations.count, *self.functional.shape[1:])
    if y.shape != expected_shape:
        raise ValueError(
            f"appended values have shape {y.shape}; expected {expected_shape}"
        )
    values = jax.lax.dynamic_update_slice_in_dim(self.y, y, self.count, axis=0)
    return ObservationBlock(design, values)

FunctionalDataset 

FunctionalDataset(blocks: Sequence[ObservationBlock])

Bases: Module

An ordered collection of observed measurement blocks.

Parameters:

Name Type Description Default
blocks Sequence[ObservationBlock]

Nonempty sequence of observation blocks with unique names and a common input dimension. Each block retains its own variance parameter.

 required
Source code in src/gpex/blocks/datasets.py 
152
153
154
155
156
def __init__(self, blocks: Sequence[ObservationBlock]) -> None:
    blocks = tuple(blocks)
    self.blocks = blocks
    _validate_collection(self.functional_blocks)
    self.index_map = IndexMap.from_blocks(self.functional_blocks)

blocks instance-attribute 

blocks: tuple[ObservationBlock, ...] = blocks

index_map class-attribute instance-attribute 

index_map: IndexMap = IndexMap.from_blocks(
    self.functional_blocks
)

functional_blocks property 

functional_blocks: tuple[FunctionalBlock, ...]

query property 

query: FunctionalQuery

design property 

design: ObservationDesign

values_flat property 

values_flat: ObservationVector

active_size property 

active_size: Count

variance_flat property 

variance_flat: ObservationVector

Measurement-error variances in canonical scalar-output order.

StructuredValues 

StructuredValues(flat: FlatValues, index_map: IndexMap)

Bases: Module

Flat and named views of means, variances, or samples.

Parameters:

Name Type Description Default
flat FlatValues

Numerical values with shape (*leading, index_map.size).

 required
index_map IndexMap

Static mapping from names to flattened slices and shapes.

 required

Attributes:

Name Type Description
flat FlatValues

Array of shape (*leading, q). Samples have a leading draw axis; means and variances have no leading axes.
by_name Mapping[str, FunctionalValues]

Views derived from flat, of shape (*leading, *block.shape). The same arrays are available as values[name].
index_map IndexMap

Static mapping from names to flattened slices and shapes.
Source code in src/gpex/blocks/indexing.py 
111
112
113
114
115
116
def __init__(self, flat: FlatValues, index_map: IndexMap) -> None:
    flat = jnp.asarray(flat)
    if flat.ndim == 0 or flat.shape[-1] != index_map.size:
        raise ValueError(f"last axis must have size {index_map.size}")
    self.flat = flat
    self.index_map = index_map

flat instance-attribute 

flat: FlatValues = flat

index_map class-attribute instance-attribute 

index_map: IndexMap = index_map

by_name property 

by_name: Mapping[str, FunctionalValues]

IndexMap dataclass 

IndexMap(blocks: tuple[BlockIndex, ...], size: int)

blocks instance-attribute 

blocks: tuple[BlockIndex, ...]

size instance-attribute 

size: int

from_blocks classmethod 

from_blocks(blocks: Sequence[FunctionalBlock]) -> IndexMap
Source code in src/gpex/blocks/indexing.py 
41
42
43
44
45
46
47
48
49
50
51
52
@classmethod
def from_blocks(
    cls,
    blocks: Sequence[FunctionalBlock],
) -> IndexMap:
    indices: list[BlockIndex] = []
    start = 0
    for block in blocks:
        stop = start + block.size
        indices.append(BlockIndex(block.name, block.shape, start, stop))
        start = stop
    return cls(tuple(indices), start)

append 

append(block: FunctionalBlock) -> IndexMap

Extend the public layout, requiring a new block name.

Source code in src/gpex/blocks/indexing.py 
54
55
56
57
58
def append(self, block: FunctionalBlock) -> IndexMap:
    """Extend the public layout, requiring a new block name."""
    stop = self.size + block.size
    index = BlockIndex(block.name, block.shape, self.size, stop)
    return IndexMap((*self.blocks, index), stop)

block 

block(name: str) -> BlockIndex
Source code in src/gpex/blocks/indexing.py 
60
61
62
63
64
def block(self, name: str) -> BlockIndex:
    for block in self.blocks:
        if block.name == name:
            return block
    raise KeyError(name)

unpack 

unpack(flat: FlatValues) -> dict[str, FunctionalValues]
Source code in src/gpex/blocks/indexing.py 
66
67
68
69
70
71
72
73
74
def unpack(self, flat: FlatValues) -> dict[str, FunctionalValues]:
    flat = jnp.asarray(flat)
    if flat.ndim == 0 or flat.shape[-1] != self.size:
        raise ValueError(f"last axis must have size {self.size}")
    leading_shape = flat.shape[:-1]
    return {
        block.name: flat[..., block.slice].reshape((*leading_shape, *block.shape))
        for block in self.blocks
    }

pack 

pack(values: Mapping[str, FunctionalValues]) -> FlatValues
Source code in src/gpex/blocks/indexing.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def pack(self, values: Mapping[str, FunctionalValues]) -> FlatValues:
    extra = set(values) - {block.name for block in self.blocks}
    if extra:
        raise ValueError(f"unexpected blocks: {sorted(extra)}")
    parts = []
    for block in self.blocks:
        value = jnp.asarray(values[block.name])
        if value.shape != block.shape:
            raise ValueError(
                f"block {block.name!r} has shape {value.shape}; "
                f"expected {block.shape}"
            )
        parts.append(value.reshape(-1))
    return jnp.concatenate(parts)

ObservationDesignBlock.variance stores the scalar measurement-error variance. Use a fixed scalar or a scalar PositiveParameter to fit it. The same variance applies to every location and output component in that block.

The block's variances property evaluates that parameter and repeats it over the functional output shape. Collections expose these measurement-error variances in canonical scalar-output order through variance_flat.

Observation designs reserve max_size locations and keep a dynamic active count. See JAX and observation capacity for capacity, padding and compilation behavior. Functional queries are unreserved.