state_space_mll#

gpjax.state_space.state_space_mll(posterior, train_data, *, observation_mask=None)[source]#

Marginal log-likelihood via the square-root Kalman filter.

Internally:

  1. Unwraps the posterior (resolving paramax-wrapped parameters).

  2. Builds the SDE via to_sde(kernel).

  3. Centres targets with y - mean_function(X).

  4. Computes sigma_eff = sqrt(obs_stddev² + prior.jitter).

  5. Delegates to kalman_filter.

Pure-JAX. Assumes time-sorted input. Unsorted times yield negative Δt and silently incorrect (NaN/garbage) results — there is no internal sort, because a data-dependent reorder inside this traced objective is avoided to keep it jit/grad/MCMC-clean. Sorting and validation are the responsibility of the eager state_space.fit* wrappers (sort_state_space_data warns and reorders; validate_state_space_data checks finiteness/shape). Callers invoking this objective directly (e.g. custom MCMC/optimisers) must pre-sort by time, e.g. via sort_state_space_data(X, y, mask).

See plans/2026-04-21-state-space-gps-design.md §Stage 1.

Example

>>> import jax.numpy as jnp
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior, state_space_mll
>>> 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
>>> train_data = gpx.Dataset(X=X, y=y)
>>> mll = state_space_mll(posterior, train_data)
>>> bool(jnp.isfinite(mll))
True

Expand for references to gpjax.state_space.state_space_mll

state_space_mll

State-Space (Markovian) Gaussian Processes / Scalability