JointModel

class gpjax.gps.JointModel(prior, likelihood)

Bases: _SummaryMixin, Module, Generic[M, K, L]

The joint distribution \(p(f, y) = p(y \mid f)\,p(f)\).

Pairs a Prior with a likelihood. This is the trainable object: gpx.fit optimises its hyperparameters. Conditioning it on data — model.condition(D) or model | D — produces the posterior process.

The base class carries no inference of its own; concrete subclasses (ConjugateModel, NonConjugateModel, HeteroscedasticModel) define what conditioning means for their likelihood. A bare JointModel is a lightweight pairing used where inference is delegated elsewhere (e.g. variational families over a latent noise process).

Parameters:

• prior (Prior)

• likelihood (Any)

condition(train_data)

Condition the joint model on data, returning the posterior process.

Parameters:

train_data (Dataset) – The observations to condition on.

Returns:

The conditioned process \(p(f \mid \mathcal{D})\).

Return type:

Posterior

predict(test_inputs, train_data, *, covariance='dense')

Sugar: condition on train_data and query at test_inputs.

Defined as exactly self.condition(train_data)(test_inputs). When making repeated predictions, condition once and reuse the returned posterior — the factorisation is cached there.

Parameters:

• test_inputs (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D']) – A Jax array of test inputs.

• train_data (Dataset) – A gpx.Dataset to condition on.

• covariance (Literal['dense', 'diagonal']) – Whether to return the dense joint covariance at the test inputs or only the marginal (diagonal) variances.

Returns:

The predictive distribution.

Return type:

GaussianDistribution