fit_lbfgs#

gpjax.fit.fit_lbfgs(*, model, objective, train_data, max_iters=100, safe=True, max_linesearch_steps=32, gtol=1e-5)[source]#

Train a Module model with respect to a supplied Objective function.

Uses Optax’s L-BFGS implementation with a jax.lax.while_loop.

Parameters:

  • model (Module) – The model to be optimised.

  • objective (Objective) – The objective function to minimise.

  • train_data (Dataset) – The training data used to evaluate the objective.

  • max_iters (int) – Maximum number of L-BFGS iterations. Defaults to 100.

  • safe (bool) – Whether to validate inputs before optimisation. Defaults to True.

  • max_linesearch_steps (int) – Maximum number of line-search steps per iteration. Defaults to 32.

  • gtol (float) – Terminate if the L2 norm of the gradient falls below this threshold. Defaults to 1e-5.

Returns:

A tuple of the optimised model and the final loss

value.

Return type:

tuple[Module, Array]

Example

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import gpjax as gpx
>>> import jax.numpy as jnp

>>> xtrain = jnp.linspace(0, 1, 20).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)

>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood

>>> nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d)
>>> trained_model, final_loss = gpx.fit_lbfgs(
...     model=posterior, objective=nmll, train_data=D
... )