fit#

gpjax.state_space.fit(*, model, train_data, optim, observation_mask=None, key=None, num_iters=100, batch_size=-1, log_rate=10, verbose=True, unroll=1, safe=True)[source]#

Fit a state-space posterior with Optax (gradient-descent style).

Thin wrapper around gpx.fit. Rejects batch_size != -1 because state-space MLL is intrinsically full-batch (the temporal scan cannot be minibatched without breaking the Markov chain).

Example

>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import optax as ox
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior, fit
>>> X = jnp.linspace(0.0, 5.0, 20).reshape(-1, 1)
>>> y = jnp.sin(X)
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> fitted, history = fit(
...     model=posterior,
...     train_data=gpx.Dataset(X=X, y=y),
...     optim=ox.adam(1e-2),
...     num_iters=2,
...     key=jr.key(0),
...     verbose=False,
... )
>>> bool(jnp.all(jnp.isfinite(history)))
True

Expand for references to gpjax.state_space.fit

fit

Parameters: