Skip to content

Fitting and objectives

See Fitting hyperparameters for an end-to-end MLE and MAP example. Fitting functions optimize explicit Parameter leaves in a model PyTree. To fit both GP and observation-noise parameters, include both objects, for example (gp, dataset), in the fitted model.

Optimizers and results

Function Use
fit_lbfgs(model, objective, ...) Deterministic objective with line search and gradient convergence
fit(model, objective, optimizer, ...) A fixed number of updates using a supplied Optax optimizer
fit_multistart(starts, fit_one) Independent fits from several initial parameter settings

Single fits return the final model, loss, gradient_norm, and update and evaluation counts. fit_lbfgs sets converged=True when it meets its gradient tolerance. Reaching its iteration limit returns converged=False. Fixed-step fit always reports converged=False.

losses includes the initial loss. L-BFGS allocates max_steps + 1 entries and fills unused entries with NaN. Outside JIT, use result.losses[:int(result.step_count) + 1] for the completed history. A failed line search raises, as does an accepted iterate with non-finite parameters, objective, or gradients.

Multistart inputs may differ in trainable parameter values. Their shapes, transforms, priors, names, and fixed data must agree. result.model selects the smallest final objective, and plural diagnostic fields report each start in input order. Any failed run raises. See the multistart example.

The fitting loops are compiled. Changing optimizer configuration, model structure, shapes, or iteration budgets may recompile them. fit supports reverse differentiation through its updates. fit_lbfgs differentiates the objective but does not support reverse differentiation through the entire fit. Call multistart helpers outside an outer JIT because best-start selection returns a Python index. See JAX and observation capacity for data shape reuse.

gpex.fitting

Explicit Optax fitting, L-BFGS, and independent multistart orchestration.

FitResult

Bases: Module

Final model and its diagnostics.

losses includes the initial objective. L-BFGS fills unused entries with NaN; use the first step_count + 1 entries when plotting. converged reports L-BFGS's gradient criterion and is always false for fixed-step fits. evaluation_count includes line-search trials. optimizer_state belongs to the final iterate. Reaching the iteration limit returns a finite model; encountering a non-finite accepted iterate raises instead.

model instance-attribute 

model: Model

losses instance-attribute 

losses: LossHistory

optimizer_state instance-attribute 

optimizer_state: Any

loss instance-attribute 

loss: Scalar

gradient_norm instance-attribute 

gradient_norm: Scalar

step_count instance-attribute 

step_count: ScalarInteger

evaluation_count instance-attribute 

evaluation_count: ScalarInteger

converged instance-attribute 

converged: ScalarBoolean

MultiStartFitResult dataclass 

MultiStartFitResult(
    models: tuple[Model, ...],
    losses: MultiStartLossHistory,
    optimizer_state: Any,
    final_losses: StartScalars,
    gradient_norms: StartScalars,
    step_counts: StartIntegers,
    evaluation_counts: StartIntegers,
    converged: StartBooleans,
    best_index: int,
)

Independent fits in input order, selected by their final objective.

final_losses holds one value per start; losses holds their histories. model and loss select the start at best_index. Optimizer state and all diagnostic arrays have a leading start dimension.

models instance-attribute 

models: tuple[Model, ...]

losses instance-attribute 

losses: MultiStartLossHistory

optimizer_state instance-attribute 

optimizer_state: Any

final_losses instance-attribute 

final_losses: StartScalars

gradient_norms instance-attribute 

gradient_norms: StartScalars

step_counts instance-attribute 

step_counts: StartIntegers

evaluation_counts instance-attribute 

evaluation_counts: StartIntegers

converged instance-attribute 

converged: StartBooleans

best_index instance-attribute 

best_index: int

model property 

model: Model

loss property 

loss: Scalar

fit 

fit(
    model: Model,
    objective: Objective[Model],
    optimizer: GradientTransformation,
    *,
    step_count: int,
) -> FitResult[Model]

Minimize a scalar objective for exactly step_count updates.

Only explicit Parameter leaves are trained. Supply an Optax optimizer such as optax.adam(0.01) or a composed transformation whose update needs only gradients, state, and parameters. Use fit_lbfgs for L-BFGS.

The result contains the final iterate, its optimizer state, and a loss trace of length step_count + 1. There is no early stopping. The compiled scan supports reverse-mode differentiation through the optimization trajectory. Every iterate is checked for finite parameters, objective, and gradient. Any failure raises after the loop completes, even if later iterates are finite.

Source code in src/gpex/fitting/optax.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def fit(
    model: Model,
    objective: Objective[Model],
    optimizer: optax.GradientTransformation,
    *,
    step_count: int,
) -> FitResult[Model]:
    """Minimize a scalar objective for exactly ``step_count`` updates.

    Only explicit ``Parameter`` leaves are trained. Supply an Optax optimizer
    such as ``optax.adam(0.01)`` or a composed transformation whose update needs
    only gradients, state, and parameters. Use ``fit_lbfgs`` for L-BFGS.

    The result contains the final iterate, its optimizer state, and a loss trace
    of length ``step_count + 1``. There is no early stopping. The compiled scan
    supports reverse-mode differentiation through the optimization trajectory.
    Every iterate is checked for finite parameters, objective, and gradient.
    Any failure raises after the loop completes, even if later iterates are finite.
    """
    validate_positive_integer("step_count", step_count)
    return _fit(model, objective, optimizer, step_count=step_count)

fit_lbfgs 

fit_lbfgs(
    model: Model,
    objective: Objective[Model],
    *,
    max_steps: int = 100,
    gradient_tolerance: float = 1e-05,
    memory_size: int = 10,
    max_linesearch_steps: int = 20,
) -> FitResult[Model]

Fit a deterministic objective with Optax L-BFGS and a zoom line search.

Stop when the gradient norm is at most gradient_tolerance or after max_steps updates. Only explicit Parameter leaves are trained. Cached line-search values and gradients avoid reevaluating accepted points. Non-finite accepted iterates and failed line searches raise an error.

The result contains the final iterate. Its loss trace has length max_steps + 1, with NaN after step_count + 1 valid entries. Check converged to distinguish gradient convergence from the iteration limit. This dynamic loop cannot be reverse-mode differentiated through; the objective itself is differentiated normally.

Source code in src/gpex/fitting/lbfgs.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def fit_lbfgs(
    model: Model,
    objective: Objective[Model],
    *,
    max_steps: int = 100,
    gradient_tolerance: float = 1e-5,
    memory_size: int = 10,
    max_linesearch_steps: int = 20,
) -> FitResult[Model]:
    """Fit a deterministic objective with Optax L-BFGS and a zoom line search.

    Stop when the gradient norm is at most ``gradient_tolerance`` or after
    ``max_steps`` updates. Only explicit ``Parameter`` leaves are trained.
    Cached line-search values and gradients avoid reevaluating accepted points.
    Non-finite accepted iterates and failed line searches raise an error.

    The result contains the final iterate. Its loss trace has length
    ``max_steps + 1``, with NaN after ``step_count + 1`` valid entries. Check
    ``converged`` to distinguish gradient convergence from the iteration limit.
    This dynamic loop cannot be reverse-mode differentiated through; the
    objective itself is differentiated normally.
    """
    validate_positive_integer("max_steps", max_steps)
    validate_positive_integer("memory_size", memory_size)
    validate_positive_integer("max_linesearch_steps", max_linesearch_steps)
    validate_tolerance(gradient_tolerance)
    return _fit_lbfgs(
        model,
        objective,
        max_steps=max_steps,
        gradient_tolerance=gradient_tolerance,
        memory_size=memory_size,
        max_linesearch_steps=max_linesearch_steps,
    )

fit_multistart 

fit_multistart(
    starts: Sequence[Model],
    fit_one: Callable[[Model], FitResult[Model]],
) -> MultiStartFitResult[Model]

Vectorize a fitting callable over informed parameter starts.

For example, fit_multistart(starts, lambda m: fit_lbfgs(m, objective)). Starts may differ only in their explicit parameter arrays: structure, shapes, parameter metadata, and fixed leaves must agree. Each run owns its optimizer state. Any failed run raises; finite final objectives determine the selected model. This function selects best_index on the host; use the single-fit functions when nesting fitting inside another JIT.

Source code in src/gpex/fitting/multistart.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def fit_multistart(
    starts: Sequence[Model],
    fit_one: Callable[[Model], FitResult[Model]],
) -> MultiStartFitResult[Model]:
    """Vectorize a fitting callable over informed parameter starts.

    For example, ``fit_multistart(starts, lambda m: fit_lbfgs(m, objective))``.
    Starts may differ only in their explicit parameter arrays: structure,
    shapes, parameter metadata, and fixed leaves must agree. Each run owns its
    optimizer state. Any failed run raises; finite final objectives determine
    the selected model. This function selects ``best_index`` on the host; use
    the single-fit functions when nesting fitting inside another JIT.
    """
    runs, models = _map_starts(starts, fit_one)
    return MultiStartFitResult(
        models=models,
        losses=runs.losses,
        optimizer_state=runs.optimizer_state,
        final_losses=runs.loss,
        gradient_norms=runs.gradient_norm,
        step_counts=runs.step_count,
        evaluation_counts=runs.evaluation_count,
        converged=runs.converged,
        best_index=int(jax.device_get(jnp.argmin(runs.loss))),
    )

Subset objectives

fit_random_subsets draws a fresh subset without replacement from each observation block at every Optax update. An integer subset_size applies to every block, or a mapping gives counts by block name. All output components at a selected location stay together, along with the block's shared noise. Each selected count must be between one and the active location count.

A GP marginal likelihood couples observations through its covariance. A subset objective is therefore an approximation, even after likelihood scaling. For MAP, scale only the likelihood and include the prior once. Use a full-data reference_objective to evaluate the final model. The subset fitting example shows the complete objective and its scaling.

subset_losses records each update's proxy objective. Values from different subsets are not directly comparable. reference_loss is evaluated once at the final model. Multistart subset fits select by their final full-data reference_losses. Use a first-order Optax optimizer for the changing subset objective.

gpex.fitting

Explicit Optax fitting, L-BFGS, and independent multistart orchestration.

RandomSubsetFitResult

Bases: Module

Final subset iterate, assessed once on the full reference objective.

subset_losses records the objective before each update. Values from different subsets are not directly comparable. reference_loss and reference_gradient_norm describe the returned model on the full data.

model instance-attribute 

model: Model

subset_losses instance-attribute 

subset_losses: LossHistory

optimizer_state instance-attribute 

optimizer_state: Any

reference_loss instance-attribute 

reference_loss: Scalar

reference_gradient_norm instance-attribute 

reference_gradient_norm: Scalar

step_count instance-attribute 

step_count: ScalarInteger

evaluation_count instance-attribute 

evaluation_count: ScalarInteger

loss property 

loss: Scalar

MultiStartRandomSubsetFitResult dataclass 

MultiStartRandomSubsetFitResult(
    models: tuple[Model, ...],
    subset_losses: MultiStartLossHistory,
    optimizer_state: Any,
    reference_losses: StartScalars,
    reference_gradient_norms: StartScalars,
    step_counts: StartIntegers,
    evaluation_counts: StartIntegers,
    best_index: int,
)

Independent subset fits, selected by their final reference objective.

models instance-attribute 

models: tuple[Model, ...]

subset_losses instance-attribute 

subset_losses: MultiStartLossHistory

optimizer_state instance-attribute 

optimizer_state: Any

reference_losses instance-attribute 

reference_losses: StartScalars

reference_gradient_norms instance-attribute 

reference_gradient_norms: StartScalars

step_counts instance-attribute 

step_counts: StartIntegers

evaluation_counts instance-attribute 

evaluation_counts: StartIntegers

best_index instance-attribute 

best_index: int

model property 

model: Model

loss property 

loss: Scalar

fit_random_subsets 

fit_random_subsets(
    model: GPModel,
    subset_objective: Objective[GPModel],
    optimizer: GradientTransformation,
    *,
    step_count: int,
    subset_size: SubsetSize,
    reference_objective: Objective[GPModel],
    key: PRNGKeyArray,
) -> RandomSubsetFitResult[GPModel]

Fit (gp, dataset) with a fresh subset per Optax update.

Supply a first-order optimizer such as optax.adam(0.01). subset_size is an active location count per block, or a mapping from every block name to its count. Sampling is uniform without replacement over active locations and reproducible from key. Every block needs at least that many active rows. A subset GP likelihood is an approximation, not an unbiased estimate of the full likelihood. The caller defines any scaling in subset_objective.

Run exactly step_count updates, then evaluate reference_objective and its gradient once on the full data. Return the final model; non-finite iterates, subset diagnostics, or reference diagnostics raise an error.

Source code in src/gpex/fitting/subsets.py 
42
43
44
45
46
47
48
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
def fit_random_subsets(
    model: GPModel,
    subset_objective: Objective[GPModel],
    optimizer: optax.GradientTransformation,
    *,
    step_count: int,
    subset_size: SubsetSize,
    reference_objective: Objective[GPModel],
    key: PRNGKeyArray,
) -> RandomSubsetFitResult[GPModel]:
    """Fit ``(gp, dataset)`` with a fresh subset per Optax update.

    Supply a first-order optimizer such as ``optax.adam(0.01)``. ``subset_size``
    is an active location count per block, or a mapping from every block name to
    its count. Sampling is uniform without replacement over active locations and
    reproducible from ``key``. Every block needs at least that many active rows.
    A subset GP likelihood is an approximation, not an unbiased estimate of the
    full likelihood. The caller defines any scaling in ``subset_objective``.

    Run exactly ``step_count`` updates, then evaluate ``reference_objective``
    and its gradient once on the full data. Return the final model; non-finite
    iterates, subset diagnostics, or reference diagnostics raise an error.
    """
    validate_positive_integer("step_count", step_count)
    _validate_gp_model(model)
    sizes = _resolve_subset_sizes(model[1], subset_size)
    return _fit_random_subsets(
        model,
        subset_objective,
        optimizer,
        step_count=step_count,
        subset_sizes=sizes,
        reference_objective=reference_objective,
        key=key,
    )

fit_multistart_random_subsets 

fit_multistart_random_subsets(
    starts: Sequence[GPModel],
    subset_objective: Objective[GPModel],
    optimizer: GradientTransformation,
    *,
    step_count: int,
    subset_size: SubsetSize,
    reference_objective: Objective[GPModel],
    key: PRNGKeyArray,
) -> MultiStartRandomSubsetFitResult[GPModel]

Run fit_random_subsets with an independent key per start.

Arguments and failure behavior match fit_random_subsets. Starts must satisfy the same structure and fixed-leaf requirements as fit_multistart. Select the final model with the lowest full-data reference objective.

Source code in src/gpex/fitting/multistart.py 
48
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
def fit_multistart_random_subsets(
    starts: Sequence[GPModel],
    subset_objective: Objective[GPModel],
    optimizer: optax.GradientTransformation,
    *,
    step_count: int,
    subset_size: SubsetSize,
    reference_objective: Objective[GPModel],
    key: PRNGKeyArray,
) -> MultiStartRandomSubsetFitResult[GPModel]:
    """Run ``fit_random_subsets`` with an independent key per start.

    Arguments and failure behavior match ``fit_random_subsets``. Starts must
    satisfy the same structure and fixed-leaf requirements as ``fit_multistart``.
    Select the final model with the lowest full-data reference objective.
    """
    keys = jax.random.split(key, len(starts))

    def fit_one(model: GPModel, start_key: PRNGKeyArray):
        return fit_random_subsets(
            model,
            subset_objective,
            optimizer,
            step_count=step_count,
            subset_size=subset_size,
            reference_objective=reference_objective,
            key=start_key,
        )

    runs, models = _map_starts(starts, fit_one, keys)
    return MultiStartRandomSubsetFitResult(
        models=models,
        subset_losses=runs.subset_losses,
        optimizer_state=runs.optimizer_state,
        reference_losses=runs.reference_loss,
        reference_gradient_norms=runs.reference_gradient_norm,
        step_counts=runs.step_count,
        evaluation_counts=runs.evaluation_count,
        best_index=int(jax.device_get(jnp.argmin(runs.reference_loss))),
    )

Exact objectives

gpex.fitting.objectives

Differentiable exact GP objectives.

negative_log_marginal_likelihood 

negative_log_marginal_likelihood(
    gp: GaussianProcess,
    dataset: FunctionalDataset,
    *,
    jitter: float = 1e-06,
) -> ScalarObjective

Return the exact Gaussian negative log marginal likelihood.

Invalid factorizations produce non-finite trials for optimizers to reject. Padding contributes neither observations nor likelihood normalization.

Source code in src/gpex/fitting/objectives.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def negative_log_marginal_likelihood(
    gp: GaussianProcess,
    dataset: FunctionalDataset,
    *,
    jitter: float = 1e-6,
) -> ScalarObjective:
    """Return the exact Gaussian negative log marginal likelihood.

    Invalid factorizations produce non-finite trials for optimizers to reject.
    Padding contributes neither observations nor likelihood normalization.
    """
    measurements = prepare_functionals(tuple(block.design for block in dataset.blocks))
    variance = dataset.variance_flat
    system = CovarianceSystem(measurements, gp.kernel, variance, jitter)
    prepared = gp.solver.prepare(
        system,
        materialization_tile_size=gp.covariance_tile_size,
        active=measurements.active,
    )
    residual = dataset.values_flat - assemble_mean(measurements, gp.mean)
    quadratic = prepared.quadratic_form(residual)
    logdet = prepared.logdet()
    normalizer = dataset.active_size * jnp.log(2 * jnp.pi)
    return 0.5 * (quadratic + logdet + normalizer)

negative_log_posterior 

negative_log_posterior(
    gp: GaussianProcess,
    dataset: FunctionalDataset,
    *,
    jitter: float = 1e-06,
) -> ScalarObjective

Add explicitly declared GP and noise priors to the exact likelihood.

Source code in src/gpex/fitting/objectives.py 
44
45
46
47
48
49
50
51
52
def negative_log_posterior(
    gp: GaussianProcess,
    dataset: FunctionalDataset,
    *,
    jitter: float = 1e-6,
) -> ScalarObjective:
    """Add explicitly declared GP and noise priors to the exact likelihood."""
    likelihood = negative_log_marginal_likelihood(gp, dataset, jitter=jitter)
    return likelihood - log_prior((gp, dataset), require_all=True)