Releases · QuantClimate/GPJax

Release list

GPJax v1.0.0 Latest

Latest

github-actions released this 28 Sep 08:39

v1.0.0

e44bb7a

GPJax 1.0 is the first stable release. It rebuilds the core API around conditioning, so the code reads like the maths: a joint model is conditioned on data to give a posterior, and the posterior is queried directly. On top of that come natural-gradient training, the dual (t-SVGP) parameterisation of sparse GPs, xarray support for gridded data, a robust Student's t likelihood, and dense joint covariances for state-space GPs.

GPJax now lives at QuantClimate/GPJax, with documentation at gpjax.quantclimate.com. Old links to docs.jaxgaussianprocesses.com redirect permanently.

Important

1.0 contains breaking changes. Work through the 0.18 → 1.0 migration guide when upgrading. Most prediction code keeps working unchanged.

✨ Highlights

Conditioning is the core API

prior * likelihood now returns a JointModel, the joint $p(f, y)$ that gpx.fit trains. Conditioning it on data gives a Posterior that caches its factorisation and is queried directly:

model = prior * gpx.likelihoods.Gaussian()
posterior = model.condition(D)          # or: model | D
predictive = posterior(xtest)           # condition once, predict many times
marginals = posterior(xtest, covariance="diagonal")
evidence = posterior.log_marginal_likelihood

One contract for every process: exact (ExactPosterior), non-conjugate (LatentPosterior), sparse variational (SparsePosterior), collapsed (CollapsedPosterior), state-space (Kalman) and OILMM posteriors all condition and predict the same way, with a single covariance= keyword.
Variational families condition through the same machinery. q.condition(D) returns a Posterior, exactly like model.condition(D).
Pathwise sampling (sample_approx) now also works on sparse posteriors.
One jitter setting: Prior.jitter is applied exactly once, inside conditioning, via the new gpjax.linalg.stabilised_cholesky. Prediction and the marginal likelihood can no longer factorise different matrices.
Likelihoods are pure conditional distributions and no longer take num_datapoints. The minibatch ELBO scale comes from the dataset itself (Dataset.n_total), so it can't be set wrong.
The design is recorded in ADR-0001, and the vocabulary in CONTEXT.md.

Natural gradients: `gpx.fit_natgrads`

Trains a sparse variational GP by alternating a natural-gradient step on $q(u)$ with an Optax step on the hyperparameters and inducing inputs, following Salimbeni et al. (2018). It supports VariationalGaussian and WhitenedVariationalGaussian. On a conjugate full-batch problem, one step with natgrad_lr=1.0 reaches the optimal $q$.

Dual parameterisation (t-SVGP): `DualVariationalGaussian` and `dual_elbo`

The dual parameterisation of Adam et al. (2021) stores Gaussian sites instead of the moments of $q(u)$. A natural-gradient step then becomes a closed-form convex combination, and the M-step gets t-SVGP's better-behaved hyperparameter gradients. fit_natgrads dispatches on the family automatically.

Gridded data with xarray: `gpjax.xarray`

from_xarray(ds, target=..., inputs=[...]) flattens a labelled xr.Dataset into an ordinary Dataset. Inputs can be coordinates or data variables, datetimes are converted to days, and NaN cells are dropped. It also returns a GridSpec, which builds inputs for any new grid and maps predictions back onto it: means and variances, or joint posterior samples with a sample dimension. Install with pip install "gpjax[xarray]"; import gpjax never imports xarray.

Robust regression: `gpx.likelihoods.StudentT`

A heavy-tailed Student's t likelihood, so outliers pull the posterior mean less strongly. Its expected log-likelihood is computed by Gauss–Hermite quadrature.

State-space GPs

covariance="dense" on StateSpacePrior.predict, StateSpaceConjugateModel.predict and the smoothed StateSpacePosterior returns the full joint covariance across test points. It is built from the RTS smoother's cross-covariance recursion, so the cost stays linear in the number of training points.
TruncatedPeriodic × Matérn product kernels now have an SDE representation (quasi-periodic models).
rts_smoother is now a genuine QR-based square-root smoother.

Parameters: `gpjax.parameters.val`

val(param) is now public and the single way to read a parameter's constrained value. Plain arrays pass through unchanged, and frozen (paramax.non_trainable) parameters are read like any other. Models are no longer unwrapped behind the scenes, so custom kernels, means, likelihoods and objectives see the same model whether you call them directly or through fit.

🚨 Breaking changes

See the migration guide for before-and-after code. In brief:

0.18	1.0
`ConjugatePosterior` / `NonConjugatePosterior`	`ConjugateModel` / `NonConjugateModel`
`HeteroscedasticPosterior` / `ChainedPosterior`	`HeteroscedasticModel` (takes `noise_prior=` directly)
`construct_posterior`	`construct_model`
`AbstractPosterior`	split into `JointModel` and `Posterior`
`StateSpaceConjugatePosterior`	`StateSpaceConjugateModel`
`posterior.predict(x, D)`	`model.condition(D)(x)` (the old call still works)
`return_covariance_type=`	`covariance=`
`Gaussian(num_datapoints=n, ...)`	`Gaussian(...)`, and likewise for every likelihood
`HeteroscedasticGaussian(noise_prior=...)`	the noise prior moves onto `HeteroscedasticModel`
variational family `posterior=`, `jitter=`	`model=`; set jitter on the `Prior`
`NaturalVariationalGaussian`, `ExpectationVariationalGaussian`	removed; use `VariationalGaussian` with `fit_natgrads`
`OILMMModel.condition_on_observations(D)`	`model.condition(D)` (the old name is deprecated, not removed)
`OILMMPosterior` default covariance `"dense"`	`"diagonal"`
custom code reading `kernel.lengthscale` directly	wrap each read in `val(...)`

NonConjugateModel also sizes its latent vector lazily on first fit; call model.init_latent(n) to use it earlier.

🐛 Fixes

The default Zero() mean no longer drifts away from zero during fitting. This was a regression introduced by the Equinox migration (#712).
gpx.kernels.RBF() and friends now type-check with no arguments, and White no longer carries an unused trainable lengthscale (#695).
Random Fourier feature kernels no longer compute the gram matrix twice.
BlockDiag and Kronecker operators correctly report unit diagonals to Lineax.
Noise-free Gaussian likelihoods (obs_stddev=0.0) work with numpyro 0.22 (#785).

📚 Documentation

The docs moved from MkDocs to Sphinx and now live at gpjax.quantclimate.com. Old MkDocs-era URLs still resolve.
New notebooks: Natural Gradients, Natural Gradients in Practice, Dual Parameterisation of Sparse GPs (t-SVGP) and Gridded Data with xarray.
A migration guide, and a Sharp Bits section explaining how the parameter system works.

🏗️ Project

The repository moved to the QuantClimate organisation.
Releases are published to PyPI with trusted publishing, so no long-lived tokens are stored.
Development dependencies moved to PEP 735 dependency groups, and GitHub Actions are pinned and audited with zizmor.

🙌 Contributors

Thanks to everyone who contributed to 1.0: @thomaspinder, @Thomas-Christie and @stephen-huan.

Full changelog: CHANGELOG.md · v0.18.0...v1.0.0

Contributors

stephen-huan, thomaspinder, and Thomas-Christie

Assets 4

GPJax v0.18.0

github-actions released this 26 Jul 19:46

v0.18.0

177dabb

Release v0.18.0

🐛 Bug Fixes

require the noise latent in heteroscedastic link_function (#706)
parameterise the spectral measure by lengthscale and dimension (#702)

📝 Other Changes

chore(release): prepare v0.18.0 (#711)
perf(variational): compute the prior KL in closed form (#708)
perf: reuse Cholesky factors in Gaussian KL divergence (#707)
build: drop the unused tensorstore runtime dependency (#705)
docs(kernels): fix RationalQuadratic and PoweredExponential formulas (#703)
ci(ruff): pin ruff to the lockfile and stop CI auto-fixing (#704)
ci(deps): bump actions/setup-node from 6 to 7 (#700)
deps(deps): bump tqdm from 4.68.4 to 4.69.0 (#701)
Fix runtime dependency metadata (#699)
Add real data examples (#696)
deps(deps): bump tqdm from 4.68.3 to 4.68.4 (#697)
Fix Matern kernel docstring formulas (#690)
deps(deps): bump absl-py from 2.4.0 to 2.5.0 (#694)
deps(deps): bump typing-extensions from 4.15.0 to 4.16.0 (#693)
docs: clarify collapsed elbo batching (#686)
Consolidate the duplicated paramax-unwrap _val helper into gpjax.parameters (#692)
Merge pull request #691 from thomaspinder/fix/skip-pages-deploy-for-forks
ci: skip docs preview deploy for fork PRs

📊 Performance & ML Improvements

See performance regression tests in this release
Model validation results available in CI artifacts

🔍 What's Changed

Full diff: v0.17.0...v0.18.0

What's Changed

ci: skip docs preview deploy for fork PRs by @thomaspinder in #691
Consolidate the duplicated paramax-unwrap _val helper into gpjax.parameters by @businessarshgoyal in #692
docs: clarify collapsed elbo batching by @vku2018 in #686
deps(deps): bump typing-extensions from 4.15.0 to 4.16.0 by @dependabot[bot] in #693
deps(deps): bump absl-py from 2.4.0 to 2.5.0 by @dependabot[bot] in #694
Fix Matérn kernel docstring formulas by @jordansilly77-stack in #690
deps(deps): bump tqdm from 4.68.3 to 4.68.4 by @dependabot[bot] in #697
Add real data examples by @thomaspinder in #696
Fix runtime dependency declarations by @ShreyanshGoyal in #699
deps(deps): bump tqdm from 4.68.4 to 4.69.0 by @dependabot[bot] in #701
ci(deps): bump actions/setup-node from 6 to 7 by @dependabot[bot] in #700
ci(ruff): pin ruff to the lockfile and stop CI auto-fixing by @thomaspinder in #704
docs(kernels): fix RationalQuadratic and PoweredExponential formulas by @thomaspinder in #703
fix(kernels): parameterise the spectral measure by lengthscale and dimension by @thomaspinder in #702
build: drop the unused tensorstore runtime dependency by @thomaspinder in #705
fix(likelihoods): require the noise latent in heteroscedastic link_function by @thomaspinder in #706
perf: reuse Cholesky factors in Gaussian KL divergence by @thomaspinder in #707
perf: use the closed-form KL in variational families by @thomaspinder in #708
chore(release): prepare v0.18.0 by @thomaspinder in #711