Skip to content

Parameters and priors

Parameter values

Only explicit Parameter leaves are fitted. Plain numbers and arrays remain fixed. Parameters store unconstrained optimization coordinates and expose the model value through .value.

Constructor Initial value Constraint
Parameter(unconstrained=...) Optimization coordinates Identity transform by default
PositiveParameter(value) Positive model value Softplus
CholeskyParameter(value) Lower-triangular model matrix Positive diagonal and zero upper triangle

PositiveParameter is used for trainable kernel variance, lengthscales, and scalar observation-noise variance. CholeskyParameter is used for a full input precision factor. A custom transform supplies forward and inverse. Pass transform.inverse(value) to the base constructor's unconstrained argument when starting from a constrained value.

See Fitting hyperparameters for a complete example.

Priors

A prior's log_prob receives the constrained model value and returns one scalar joint log density. Scalar prior families sum independent contributions for array-valued parameters. MultivariateLogNormalPrior describes the joint normal distribution of the log-values. A correlated prior on ARD lengthscales still uses an axis-aligned kernel metric.

MAP requires a prior on every trainable parameter. MLE ignores attached priors. The optimization transform adds no Jacobian term to the MAP objective, which is defined on the constrained model values.

For a full precision factor L, IndependentCholeskyPrior uses separate priors for positive diagonal and unconstrained strict-lower entries. ScaledLKJCholeskyPrior describes L = diag(scales) @ C, with positive scales and a correlation Cholesky factor C. These correlations describe the input precision metric. The prior includes the density correction from scales and correlation coordinates to the free entries of L.

Custom priors provide a scalar log_prob(value). Built-in transforms and priors implement the Transform and Prior interfaces below.

Classes and functions

gpex.parameters

Trainable parameters, optimization transforms, and priors.

Transform

Bases: Protocol

Map between optimization coordinates and constrained parameter values.

forward abstractmethod 

forward(unconstrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
19
20
@abstractmethod
def forward(self, unconstrained: ParameterValue) -> ParameterValue: ...

inverse abstractmethod 

inverse(constrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
22
23
@abstractmethod
def inverse(self, constrained: ParameterValue) -> ParameterValue: ...

Parameter 

Parameter(
    *,
    unconstrained: ParameterValue | float,
    transform: Transform | None = None,
    prior: Prior | None = None,
    name: str | None = None,
)

Bases: Module

Trainable optimization coordinates with a transform and optional prior.

Parameters:

Name Type Description Default
unconstrained ParameterValue | float

Initial raw scalar or array, stored directly. The model value is transform.forward(unconstrained) and may have a different shape, for example a matrix represented by packed coordinates.

 required
transform Transform | None

Map from optimization coordinates to constrained values. Defaults to IdentityTransform.

 None
prior Prior | None

Distribution with log_prob(value) returning one scalar joint log density. Defaults to None; MAP requires an explicit prior.

 None
name str | None

Optional diagnostic name. Defaults to None.

 None

Construction supports JIT. Use PositiveParameter(value) or CholeskyParameter(value) to initialize from a constrained model value.

Source code in src/gpex/parameters/parameters.py 
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(
    self,
    *,
    unconstrained: ParameterValue | float,
    transform: Transform | None = None,
    prior: Prior | None = None,
    name: str | None = None,
) -> None:
    self.unconstrained = jnp.asarray(unconstrained)
    self.transform = IdentityTransform() if transform is None else transform
    self.prior = prior
    self.name = name

unconstrained instance-attribute 

unconstrained: ParameterValue = jnp.asarray(unconstrained)

transform class-attribute instance-attribute 

transform: Transform = (
    IdentityTransform() if transform is None else transform
)

prior class-attribute instance-attribute 

prior: Prior | None = prior

name class-attribute instance-attribute 

name: str | None = name

value property 

value: ParameterValue

log_prior 

log_prior() -> Scalar

Return the joint log density on value, or zero without a prior.

Source code in src/gpex/parameters/parameters.py 
65
66
67
68
69
def log_prior(self) -> Scalar:
    """Return the joint log density on ``value``, or zero without a prior."""
    if self.prior is None:
        return jnp.zeros(())
    return self.prior.log_prob(self.value)

PositiveParameter 

PositiveParameter(
    value: ParameterValue | float,
    *,
    prior: Prior | None = None,
    name: str | None = None,
)

Bases: Parameter

A positive scalar or array optimized through a softplus transform.

Parameters:

Name Type Description Default
value ParameterValue | float

Positive initial scalar or array, for example a lengthscale vector of shape (d,). Values must be finite.

 required
prior Prior | None

Optional joint prior on the complete positive value. Defaults to None. The optimization transform does not add a density Jacobian.

 None
name str | None

Optional diagnostic name. Defaults to None.

 None

Construction supports JIT, batching, and differentiation of valid values. The transform, prior, and name remain static metadata.

Source code in src/gpex/parameters/parameters.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def __init__(
    self,
    value: ParameterValue | float,
    *,
    prior: Prior | None = None,
    name: str | None = None,
) -> None:
    value = jnp.asarray(value)
    value = error_if(
        value,
        jnp.any(~jnp.isfinite(value) | (value <= 0)),
        "positive parameters must be initialized with finite values above zero",
    )
    transform = SoftplusTransform()
    super().__init__(
        unconstrained=transform.inverse(value),
        transform=transform,
        prior=prior,
        name=name,
    )

CholeskyParameter 

CholeskyParameter(
    value: ParameterValue,
    *,
    prior: Prior | None = None,
    name: str | None = None,
)

Bases: Parameter

A trainable lower-triangular matrix with positive diagonal.

Parameters:

Name Type Description Default
value ParameterValue

Complete initial matrix of shape (d, d). It must be finite, lower triangular, and have a positive diagonal.

 required
prior Prior | None

Optional joint prior on the complete constrained matrix. MAP requires a prior; maximum likelihood does not. Defaults to None.

 None
name str | None

Optional diagnostic name. Defaults to None.

 None

Only the d * (d + 1) / 2 free entries are optimized, in row-major order. Diagonal entries use softplus; strict-lower entries are unrestricted. Priors are evaluated on value without an optimization-coordinate Jacobian. Construction supports JIT, batching, and differentiation of valid values; the matrix shape, prior, and name remain static.

Source code in src/gpex/parameters/parameters.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def __init__(
    self,
    value: ParameterValue,
    *,
    prior: Prior | None = None,
    name: str | None = None,
) -> None:
    value = validate_cholesky_factor(value)
    transform = CholeskyTransform(value.shape[0])
    super().__init__(
        unconstrained=transform.inverse(value),
        transform=transform,
        prior=prior,
        name=name,
    )

IdentityTransform dataclass 

IdentityTransform()

Bases: Transform

forward 

forward(unconstrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
29
30
def forward(self, unconstrained: ParameterValue) -> ParameterValue:
    return unconstrained

inverse 

inverse(constrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
32
33
def inverse(self, constrained: ParameterValue) -> ParameterValue:
    return constrained

ExpTransform dataclass 

ExpTransform()

Bases: Transform

forward 

forward(unconstrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
39
40
def forward(self, unconstrained: ParameterValue) -> ParameterValue:
    return jnp.exp(unconstrained)

inverse 

inverse(constrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
42
43
def inverse(self, constrained: ParameterValue) -> ParameterValue:
    return jnp.log(constrained)

SoftplusTransform dataclass 

SoftplusTransform()

Bases: Transform

forward 

forward(unconstrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
49
50
def forward(self, unconstrained: ParameterValue) -> ParameterValue:
    return jax.nn.softplus(unconstrained)

inverse 

inverse(constrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
52
53
def inverse(self, constrained: ParameterValue) -> ParameterValue:
    return constrained + jnp.log(-jnp.expm1(-constrained))

CholeskyTransform dataclass 

CholeskyTransform(dimension: int)

Bases: Transform

Pack a lower-triangular factor using unconstrained diagonal coordinates.

Parameters:

Name Type Description Default
dimension int

Positive matrix dimension d, fixed under JAX transformations.

 required

forward maps a vector of shape (d * (d + 1) // 2,) to a matrix of shape (d, d). Entries follow row-major lower-triangle order. Diagonal coordinates use softplus, and strict-lower entries remain unrestricted. inverse expects a finite lower-triangular matrix with positive diagonal. Both operations support JIT, vmap, and differentiation of valid inputs.

dimension instance-attribute 

dimension: int

forward 

forward(unconstrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
73
74
75
76
77
78
79
80
81
82
83
def forward(self, unconstrained: ParameterValue) -> ParameterValue:
    rows, columns = jnp.tril_indices(self.dimension)
    # Transform only the diagonal: unrestricted entries remain untouched.
    diagonal = jnp.arange(self.dimension)
    diagonal_positions = (diagonal + 1) * (diagonal + 2) // 2 - 1
    entries = unconstrained.at[diagonal_positions].set(
        jax.nn.softplus(unconstrained[diagonal_positions])
    )
    return (
        jnp.zeros((self.dimension, self.dimension)).at[rows, columns].set(entries)
    )

inverse 

inverse(constrained: ParameterValue) -> ParameterValue
Source code in src/gpex/parameters/transforms.py 
85
86
87
88
89
90
91
def inverse(self, constrained: ParameterValue) -> ParameterValue:
    diagonal = jnp.diag(constrained)
    raw_diagonal = SoftplusTransform().inverse(diagonal)
    indices = jnp.arange(self.dimension)
    return constrained.at[indices, indices].set(raw_diagonal)[
        jnp.tril_indices(self.dimension)
    ]

log_prior 

log_prior(
    tree: Any, *, require_all: bool = False
) -> Scalar

Sum priors explicitly attached to all parameters in a PyTree.

Source code in src/gpex/parameters/parameters.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def log_prior(tree: Any, *, require_all: bool = False) -> Scalar:
    """Sum priors explicitly attached to all parameters in a PyTree."""
    parameters = jax.tree_util.tree_leaves(
        tree, is_leaf=lambda item: isinstance(item, Parameter)
    )
    parameters = [item for item in parameters if isinstance(item, Parameter)]
    if not parameters:
        return jnp.zeros(())
    if require_all:
        missing = [
            parameter.name or "<unnamed>"
            for parameter in parameters
            if parameter.prior is None
        ]
        if missing:
            raise ValueError(
                "MAP requires explicit priors for every Parameter; missing: "
                + ", ".join(missing)
            )
    return sum((parameter.log_prior() for parameter in parameters), jnp.zeros(()))

gpex.parameters.priors

Distributions on constrained parameter values.

Prior

Bases: Protocol

A scalar log density evaluated jointly on a parameter's complete value.

log_prob abstractmethod 

log_prob(value: ParameterValue) -> Scalar
Source code in src/gpex/parameters/priors.py 
25
26
@abstractmethod
def log_prob(self, value: ParameterValue) -> Scalar: ...

NormalPrior dataclass 

NormalPrior(
    mean: float = 0.0, standard_deviation: float = 1.0
)

Bases: Prior

mean class-attribute instance-attribute 

mean: float = 0.0

standard_deviation class-attribute instance-attribute 

standard_deviation: float = 1.0

log_prob 

log_prob(value: ParameterValue) -> Scalar
Source code in src/gpex/parameters/priors.py 
38
39
40
41
def log_prob(self, value: ParameterValue) -> Scalar:
    standardized = (value - self.mean) / self.standard_deviation
    normalization = jnp.log(self.standard_deviation) + 0.5 * jnp.log(2 * jnp.pi)
    return jnp.sum(-0.5 * standardized**2 - normalization)

LogNormalPrior dataclass 

LogNormalPrior(
    log_mean: float = 0.0,
    log_standard_deviation: float = 1.0,
)

Bases: Prior

Independent log-normal densities with shared scalar hyperparameters.

Parameters:

Name Type Description Default
log_mean float

Mean of each log-value. Defaults to 0.

 0.0
log_standard_deviation float

Positive standard deviation of each log-value. Defaults to 1. Vector inputs use the same distribution per entry.

 1.0

log_mean class-attribute instance-attribute 

log_mean: float = 0.0

log_standard_deviation class-attribute instance-attribute 

log_standard_deviation: float = 1.0

log_prob 

log_prob(value: ParameterValue) -> Scalar
Source code in src/gpex/parameters/priors.py 
61
62
63
64
65
66
67
68
69
def log_prob(self, value: ParameterValue) -> Scalar:
    log_value = jnp.log(value)
    standardized = (log_value - self.log_mean) / self.log_standard_deviation
    normalization = (
        jnp.log(value)
        + jnp.log(self.log_standard_deviation)
        + 0.5 * jnp.log(2 * jnp.pi)
    )
    return jnp.sum(-0.5 * standardized**2 - normalization)

MultivariateLogNormalPrior dataclass 

MultivariateLogNormalPrior(
    log_mean: ArrayLike, log_covariance: ArrayLike
)

Bases: Prior

A joint log-normal prior on a positive vector of shape (d,).

Parameters:

Name Type Description Default
log_mean ArrayLike

Finite mean vector of log(value), with shape (d,).

 required
log_covariance ArrayLike

Finite, symmetric positive-definite covariance of log(value), with shape (d, d). A diagonal matrix gives independent entries; off-diagonal entries allow correlations.

 required

Construct the prior outside JIT. Inputs may be arrays or Python sequences; they are stored as immutable tuples because priors are static metadata on a Parameter. The log density includes the log-normal Jacobian and is evaluated on positive values, independently of their optimization transform.

Source code in src/gpex/parameters/priors.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def __init__(
    self,
    log_mean: ArrayLike,
    log_covariance: ArrayLike,
) -> None:
    mean = np.asarray(log_mean)
    covariance = np.asarray(log_covariance)
    # Conversion to Python floats below would discard NumPy imaginary parts.
    if np.iscomplexobj(mean) or np.iscomplexobj(covariance):
        raise ValueError("log mean and covariance must be real")
    if mean.ndim != 1 or mean.size == 0 or not np.all(np.isfinite(mean)):
        raise ValueError("log mean must be a nonempty finite vector")
    if covariance.shape != (mean.size, mean.size):
        raise ValueError("log covariance must have shape (d, d) matching log mean")
    if not np.all(np.isfinite(covariance)) or not np.allclose(
        covariance, covariance.T, rtol=1e-7, atol=1e-12
    ):
        raise ValueError("log covariance must be finite and symmetric")
    covariance = 0.5 * (covariance + covariance.T)
    try:
        factor = np.linalg.cholesky(covariance)
    except np.linalg.LinAlgError as error:
        raise ValueError("log covariance must be positive definite") from error
    object.__setattr__(self, "log_mean", tuple(float(x) for x in mean))
    object.__setattr__(
        self,
        "log_covariance",
        tuple(tuple(float(x) for x in row) for row in covariance),
    )
    object.__setattr__(
        self, "_scale_tril", tuple(tuple(float(x) for x in row) for row in factor)
    )

log_mean instance-attribute 

log_mean: tuple[float, ...]

log_covariance instance-attribute 

log_covariance: tuple[tuple[float, ...], ...]

log_prob 

log_prob(value: ParameterValue) -> Scalar

Return the scalar joint log density, or -inf outside support.

Parameters:

Name Type Description Default
value ParameterValue

Positive finite vector with shape (d,).

 required

Returns:

Type Description
Scalar

Scalar log density with respect to the positive vector. Entries that are nonpositive or nonfinite give a density of -inf.
Source code in src/gpex/parameters/priors.py 
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
151
def log_prob(self, value: ParameterValue) -> Scalar:
    """Return the scalar joint log density, or ``-inf`` outside support.

    Args:
        value: Positive finite vector with shape ``(d,)``.

    Returns:
        Scalar log density with respect to the positive vector. Entries that are
            nonpositive or nonfinite give a density of ``-inf``.
    """
    value = jnp.asarray(value)
    if value.shape != (len(self.log_mean),):
        raise ValueError("prior value must have shape (d,) matching log mean")
    if jnp.issubdtype(value.dtype, jnp.complexfloating):
        raise ValueError("prior value must be real")
    valid = (value > 0) & jnp.isfinite(value)
    log_value = jnp.log(jnp.where(valid, value, 1.0))
    mean = jnp.asarray(self.log_mean)
    factor = jnp.asarray(self._scale_tril)
    standardized = jsp.linalg.solve_triangular(factor, log_value - mean, lower=True)
    log_density = (
        -0.5 * jnp.sum(standardized**2)
        - jnp.sum(jnp.log(jnp.diag(factor)))
        - 0.5 * value.size * jnp.log(2 * jnp.pi)
        - jnp.sum(log_value)
    )
    return jnp.where(jnp.all(valid), log_density, -jnp.inf)

GammaPrior dataclass 

GammaPrior(concentration: float, rate: float)

Bases: Prior

concentration instance-attribute 

concentration: float

rate instance-attribute 

rate: float

log_prob 

log_prob(value: ParameterValue) -> Scalar
Source code in src/gpex/parameters/priors.py 
163
164
165
166
167
168
169
170
def log_prob(self, value: ParameterValue) -> Scalar:
    elementwise = (
        self.concentration * jnp.log(self.rate)
        - jsp.special.gammaln(self.concentration)
        + (self.concentration - 1) * jnp.log(value)
        - self.rate * value
    )
    return jnp.sum(elementwise)

IndependentCholeskyPrior dataclass 

IndependentCholeskyPrior(
    diagonal_prior: LogNormalPrior = LogNormalPrior(),
    off_diagonal_prior: NormalPrior = NormalPrior(),
)

Bases: Prior

Independent log-normal diagonal and normal strict-lower entries.

Parameters:

Name Type Description Default
diagonal_prior LogNormalPrior

Log-normal distribution shared by the positive diagonal entries. Defaults to LogNormalPrior().

 LogNormalPrior()
off_diagonal_prior NormalPrior

Normal distribution shared by the unrestricted strict-lower entries. Defaults to NormalPrior().

 NormalPrior()

The normalized joint density is with respect to the free entries of L.

diagonal_prior class-attribute instance-attribute 

diagonal_prior: LogNormalPrior = LogNormalPrior()

off_diagonal_prior class-attribute instance-attribute 

off_diagonal_prior: NormalPrior = NormalPrior()

log_prob 

log_prob(value: ParameterValue) -> Scalar

Return the joint log density, or -inf outside Cholesky support.

Source code in src/gpex/parameters/priors.py 
189
190
191
192
193
194
195
196
197
198
def log_prob(self, value: ParameterValue) -> Scalar:
    """Return the joint log density, or ``-inf`` outside Cholesky support."""
    value = validate_factor_shape(value)
    valid = is_cholesky_factor(value)
    matrix = jnp.where(valid, value, jnp.eye(value.shape[0]))
    density = self.diagonal_prior.log_prob(jnp.diag(matrix))
    density += self.off_diagonal_prior.log_prob(
        matrix[jnp.tril_indices(value.shape[0], k=-1)]
    )
    return jnp.where(valid, density, -jnp.inf)

ScaledLKJCholeskyPrior dataclass 

ScaledLKJCholeskyPrior(
    concentration: float = 1.0,
    scale_prior: Prior = LogNormalPrior(),
)

Bases: Prior

Positive row scales and an LKJ correlation factor as a prior on full L.

Parameters:

Name Type Description Default
concentration float

Positive LKJ concentration. Defaults to 1 (uniform over correlation matrices). Larger values favor weaker correlations.

 1.0
scale_prior Prior

Joint prior on the positive row-norm vector of shape (d,). Defaults to independent LogNormalPrior() densities. Custom priors must be supported on positive vectors.

 LogNormalPrior()

Decompose L = diag(scales) @ C, where each row of C has unit norm. The normalized density includes the Jacobian from (scales, C) to the free entries of L. For generalized kernels this describes the precision matrix L @ L.T, not a covariance or lengthscale matrix.

concentration class-attribute instance-attribute 

concentration: float = 1.0

scale_prior class-attribute instance-attribute 

scale_prior: Prior = LogNormalPrior()

log_prob 

log_prob(value: ParameterValue) -> Scalar

Return the joint log density, or -inf outside Cholesky support.

Source code in src/gpex/parameters/priors.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
def log_prob(self, value: ParameterValue) -> Scalar:
    """Return the joint log density, or ``-inf`` outside Cholesky support."""
    value = validate_factor_shape(value)
    valid = is_cholesky_factor(value)
    dimension = value.shape[0]
    matrix = jnp.where(valid, value, jnp.eye(dimension))
    # Scale before squaring to avoid overflow for finite large factors.
    row_max = jnp.max(jnp.abs(matrix), axis=1)
    scales = row_max * jnp.sqrt(jnp.sum((matrix / row_max[:, None]) ** 2, axis=1))
    log_scales = jnp.log(scales)
    log_diagonal = jnp.log(jnp.diag(matrix)) - log_scales

    # Row i (zero-based) has i free correlation coordinates in a unit ball.
    # Integrating its density gives pi**(i/2) * Gamma(a) / Gamma(a+i/2).
    rows = jnp.arange(1.0, dimension)
    alpha = self.concentration + (dimension - 1 - rows) / 2
    log_normalizer = jnp.sum(
        rows / 2 * jnp.log(jnp.pi)
        + jsp.special.gammaln(alpha)
        - jsp.special.gammaln(alpha + rows / 2)
    )
    lkj_density = jnp.sum((2 * alpha - 2) * log_diagonal[1:]) - log_normalizer
    # |d L_i / d (scale_i, C_i)| = scale_i**i / C_ii.
    log_jacobian = jnp.sum(rows * log_scales[1:] - log_diagonal[1:])
    density = self.scale_prior.log_prob(scales) + lkj_density - log_jacobian
    return jnp.where(valid, density, -jnp.inf)