V1.0

CHANGELOG.md

@@ -8,8 +8,122 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **Dense joint predictive covariance for state-space GPs.**
+  `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and
+  `StateSpacePosterior.__call__` (`gpjax.state_space`) now accept
+  `covariance="dense"`, returning the full joint covariance across test
+  points rather than marginal variances only
+  ([#651](https://github.com/thomaspinder/GPJax/issues/651)). For the
+  unconditioned prior this is just the kernel's own dense gram — the
+  state-space SDE is an exact representation with no training data to
+  marginalise out. For the conditioned (smoothed) posterior it is built from
+  the RTS smoother's cross-covariance recursion (Särkkä & Solin 2019 §12.2):
+  `rts_smoother` gains an opt-in `return_gains=True` that exposes its
+  already-computed per-step smoother gains, and a new
+  `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains them
+  into the `M x M` test-point covariance. This keeps the state-space
+  formulation's linear-in-`N` cost — no `N x N` gram over the training set is
+  ever formed — with the cross-covariance work landing at `O(M^2 d^3)`,
+  independent of `N` and no larger than the `O(M^2)` already required to
+  store the dense output. `StateSpacePosterior.filtered` /
+  `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep
+  raising `NotImplementedError` for `covariance="dense"`: each of their test
+  points conditions on a different information set (observations up to its
+  own timestamp), so a "joint" filtered covariance is not the dense
+  conjugate-predictive-shaped object the smoothed path now matches, and
+  extending it is left to a future issue if there is demand for it.
+
+- **`gpjax.fit_natgrads` and the `gpjax.natural_gradients` module.** Trains a variational
+  family by alternating one natural-gradient step on the variational distribution with
+  one step of a supplied Optax optimiser on everything else — kernel and likelihood
+  hyperparameters, the mean function, and the inducing inputs — following Salimbeni,
+  Eleftheriadis and Hensman (2018), [arXiv:1803.09151](https://arxiv.org/abs/1803.09151).
+  `VariationalGaussian` and `WhitenedVariationalGaussian` are supported. For a conjugate
+  model on the full batch, `natgrad_lr=1.0` reaches the optimal `q` in one iteration.
+
+- **`DualVariationalGaussian` and `gpjax.objectives.dual_elbo`.** The dual (t-SVGP)
+  parameterisation of Adam, Chang, Khan and Solin (2021),
+  [arXiv:2111.03412](https://arxiv.org/abs/2111.03412). Instead of the moments of
+  `q(u)`, the family stores an unnormalised Gaussian *site* on the centred inducing
+  outputs — `dual_vector` is the site's first natural parameter and `dual_matrix` its
+  precision — from which `q(u)` is recovered through the working matrix
+  `R = Kzz + Kzz Lambda_2 Kzz`. Because the stored coordinates are an affine image of
+  the natural parameters, a natural-gradient step is a convex combination of the
+  current sites with a closed-form target, so no expectation-to-natural round trip is
+  needed and the KL is never differentiated. `fit_natgrads` dispatches on the family
+  and takes that step; the step size means the same thing in both branches, and from
+  the same starting `q` the dual and Salimbeni E-steps produce identical iterates —
+  provided the dual branch's computed per-point curvature stays non-negative, so that
+  its `beta_floor` never engages. That holds for a genuinely log-concave likelihood;
+  GPJax's `inv_probit` clips its probabilities away from 0 and 1, which breaks
+  log-concavity in the far tails, and there the two branches diverge. The
+  dual branch restricts the step size to the interval from zero to one, since the
+  update is a convex combination.
+  `dual_elbo` has the same *value* as `elbo` at the implied moments, for any sites and
+  any hyperparameters, but a different *hyperparameter gradient from `elbo` evaluated on
+  the matched `VariationalGaussian`*: the prior part of `q` tracks the kernel while the
+  data-dependent sites stay frozen. The difference is between the two
+  parameterisations, not between the two functions — calling `elbo` directly on a
+  `DualVariationalGaussian` returns the same value and the same gradients as
+  `dual_elbo`, which is simply the batched-marginals fast path for that family. That
+  frozen-site gradient is what gives the M-step its reported behaviour, so `Kzz` must
+  not be detached and the implied moments must not be cached on the family.
+  `DualVariationalGaussian` also works with plain `gpjax.fit`, where it is ordinary
+  gradient descent in the dual coordinates. The `VariationalParametrisationSuite` ASV
+  benchmark gains a `dual` axis value.
+
+### Changed (breaking)
+
+- **`OILMMPosterior.__call__`/`.predict` default to `covariance="diagonal"`**
+  (was `"dense"`) ([#682](https://github.com/JaxGaussianProcesses/GPJax/issues/682)).
+  The dense joint covariance costs `O(m n^2 p^2)` — an `np x np` matrix built
+  from `m` dense `n x n` latent covariances — which forfeits the
+  `O(mn^3 + nmp)` scaling that OILMM exists to provide, and is unaffordable
+  well before the mean/marginal-variance query is. Marginal variances are the
+  common case and are unaffected by the mixing matrix's off-diagonal
+  structure, so the cheap path is now the default; pass
+  `covariance="dense"` explicitly for the joint covariance.
+
+### Removed
+
+- **`NaturalVariationalGaussian` and `ExpectationVariationalGaussian`.** These were
+  parameterisation-only classes with no optimiser attached: they stored the natural or
+  expectation coordinates of `q(u)` but offered no way to take a natural-gradient step
+  in them. Natural-gradient geometry belongs to the optimiser — the Fisher matrix *is*
+  the Jacobian dη/dθ, so a natural-gradient step in the natural parameters θ is exactly
+  an ordinary gradient step in the expectation parameters η, and either coordinate system
+  can be recovered on the fly from whatever the family happens to store. `fit_natgrads`
+  therefore operates directly on `VariationalGaussian` and `WhitenedVariationalGaussian`,
+  which store constraint-respecting coordinates. Users of the removed classes should
+  switch to `VariationalGaussian` with `gpjax.fit_natgrads`.
+  The `VariationalParametrisationSuite` ASV benchmark loses its `natural` and
+  `expectation` axis values; previously recorded results for those two arms are orphaned.
+
 ### Fixed
 
+- **`gpx.kernels.RBF()` was a type error, and `White()` carried a phantom
+  trainable lengthscale**
+  ([#695](https://github.com/JaxGaussianProcesses/GPJax/issues/695)). Pyright
+  synthesises `__init__` signatures from dataclass fields for any kernel that
+  inherits its `__init__` (e.g. `RBF`, `Matern12/32/52`), and those fields had
+  no defaults, so the canonical `RBF()` call was flagged as missing arguments
+  while the nonsensical `RBF(name="xyz")` type-checked cleanly (raising
+  `TypeError` at runtime). Kernel fields now carry real defaults matching
+  their `__init__` defaults, and `name` is a `ClassVar` rather than a
+  dataclass field, so the synthesised and hand-written signatures agree.
+  Separately, `White` hardcoded `lengthscale=1.0` into
+  `StationaryKernel.__init__` even though `White.__call__` never reads it,
+  so every `White` kernel carried a real, trainable `PositiveReal` leaf with
+  zero gradient that showed up in optimiser state and MCMC traces; `White`
+  now has its own minimal `__init__` and no longer carries a lengthscale at
+  all (`White().lengthscale is None`, and it is absent from
+  `jax.tree_util.tree_flatten`). The stale `_compute_base_init` workaround in
+  `kernels/base.py`, whose docstring claimed "equinox modules are frozen
+  after `super().__init__()`" -- no longer true under the pinned Equinox
+  version -- was removed in favour of a plain `super().__init__(...)` call.
+
 - **`Zero` mean function is trainable and drifts away from zero.** Fitting a
   model with the default `Zero()` mean function moved its constant towards the
   data mean (0.0 → 5.09 on a dataset with mean 5), silently changing the

CLAUDE.md

@@ -78,17 +78,18 @@ Functions `(model, Dataset) -> scalar`:
 - `conjugate_mll` / `conjugate_loocv` -- for `ConjugatePosterior`
 - `log_posterior_density` (alias `non_conjugate_mll`) -- for `NonConjugatePosterior`
 - `elbo` / `collapsed_elbo` -- for variational families
+- `dual_elbo` -- for `DualVariationalGaussian` (t-SVGP; same value as `elbo`, different hyperparameter gradient)
 - `heteroscedastic_elbo` -- for heteroscedastic models
 
 Optimise by negating: `nmll = lambda p, d: -conjugate_mll(p, d)`
 
 ### Fitting (`gpjax/fit.py`)
 
-Three optimisers: `fit()` (Optax gradient descent with scan), `fit_scipy()` (SciPy L-BFGS-B), `fit_lbfgs()` (Optax L-BFGS with `while_loop`). All handle the constrained/unconstrained bijection automatically: `paramax.unwrap(model)` is called inside the loss function, and `eqx.partition`/`eqx.combine` with `eqx.is_array` manage trainable vs static parts.
+Four optimisers: `fit()` (Optax gradient descent with scan), `fit_scipy()` (SciPy L-BFGS-B), `fit_lbfgs()` (Optax L-BFGS with `while_loop`), `fit_natgrads()` (natural-gradient steps on a variational family, alternated with Optax steps on the hyperparameters). All handle the constrained/unconstrained bijection automatically: `paramax.unwrap(model)` is called inside the loss function, and `eqx.partition`/`eqx.combine` with `eqx.is_array` manage trainable vs static parts.
 
 ### Variational inference (`gpjax/variational_families.py`)
 
-`VariationalGaussian`, `WhitenedVariationalGaussian`, `NaturalVariationalGaussian`, `ExpectationVariationalGaussian`, `CollapsedVariationalGaussian`, `GraphVariationalGaussian`, `HeteroscedasticVariationalFamily`. All inherit from `AbstractVariationalFamily` and implement `predict()` + `prior_kl()`.
+`VariationalGaussian`, `WhitenedVariationalGaussian`, `DualVariationalGaussian`, `CollapsedVariationalGaussian`, `GraphVariationalGaussian`, `HeteroscedasticVariationalFamily`. All inherit from `AbstractVariationalFamily` and implement `predict()` + `prior_kl()`.
 
 ### NumPyro integration (`gpjax/numpyro_extras.py`)
 

CONTEXT.md

@@ -0,0 +1,89 @@
+# GPJax domain glossary
+
+The ubiquitous language of this codebase. Code, docs, tests, and reviews use
+these terms exactly; when a concept is missing, add it here in the same PR
+that introduces it.
+
+## Design principle
+
+**Maths-first, comprehensible to the maths-familiar non-expert.** API names
+follow the textbook equation unless the textbook word is jargon a
+practitioner would not own. When the two conflict, choose the word a
+scikit-learn/PyMC user already speaks and state the maths once in the
+docstring (e.g. the class is `JointModel`, and its docstring says "the joint
+distribution p(f, y)").
+
+## Terms
+
+**Prior** — the Gaussian process prior p(f), pairing a kernel with a mean
+function. Queryable at any inputs: `prior(x)` returns the prior predictive.
+Owns the model's single numerical-stabilisation knob, `jitter`.
+
+**Likelihood** — the conditional distribution p(y | f). A *pure* conditional:
+it holds no priors, no dataset facts, and no data sizes. (Both were tried and
+removed at v1.0 — see ADR-0001.)
+
+**JointModel** — the joint distribution p(f, y) = p(y | f) · p(f), created by
+`prior * likelihood`. The *trainable* object: `gpx.fit` optimises its
+hyperparameters. It carries state that must exist before conditioning
+(hyperparameters; the non-conjugate latent; the heteroscedastic noise-process
+model) and no derived quantities. Concrete kinds: `ConjugateModel`,
+`NonConjugateModel`, `HeteroscedasticModel` — each holds state the others
+cannot.
+
+**condition** — the operation p(f, y) + 𝒟 → p(f | 𝒟), spelled
+`model.condition(D)` or the operator form `model | D` (read: "f given D").
+
+The signature is `condition(train_data)` uniformly, on every conditionable
+object, with `train_data` required. It is universal across the exact, latent,
+sparse, collapsed and state-space modes. Where the maths does not consume the
+data — the uncollapsed variational families, which already carry the fitted
+q(u) — the argument is still accepted, for interface uniformity, and the
+docstring says so plainly. A signature that varied by object would be a
+worse API than one argument occasionally ignored.
+
+The one documented exclusion is the heteroscedastic path
+(`HeteroscedasticModel` and `HeteroscedasticVariationalFamily`), which has no
+closed-form conditioned process: it carries two latent processes, signal and
+noise, so there is no single p(f | 𝒟) to return. Both raise
+`NotImplementedError` naming the alternative — inference runs through
+`HeteroscedasticVariationalFamily` and the `heteroscedastic_elbo` objective,
+and prediction through `predict` / `predict_latents`, or by conditioning the
+`signal_variational` and `noise_variational` components individually.
+
+`prior_kl` is deliberately *not* part of this contract: it keeps a per-family
+signature, because only the collapsed family's KL is a function of the data.
+
+**Posterior** — the conditioned process p(f | 𝒟) returned by `condition`. An
+*immutable* pytree: the training-covariance factorisation is computed once and
+cached; every query is a view of it. The uniform query surface:
+`posterior(xtest, covariance="dense"|"diagonal")`, plus per-mode views —
+`log_marginal_likelihood` / `loo` / `sample_approx` on the exact mode,
+`log_posterior_density` on the latent mode. Users never name the concrete
+implementations behind the interface.
+
+**variational family** — a trainable approximate posterior over inducing
+values: to sparse GPs what JointModel is to exact ones. It carries the joint
+model in its `model` field, and `.condition(D)` yields a Posterior like any
+other; `elbo`-style objectives are its training criteria. The model's
+`Prior.jitter` is the only stabilisation knob — families carry none of their
+own.
+
+**evidence / log marginal likelihood** — p(𝒟), the normalising constant of
+conditioning, exposed as `posterior.log_marginal_likelihood`. "Evidence" and
+"marginal likelihood" are the same quantity; the attribute uses the
+GP-community's term.
+
+**sugar** — a documented one-line composition kept for ergonomics, never a
+second implementation. `model.predict(x, D)` and `model(x, D)` are sugar for
+`model.condition(D)(x)`; `model | D` is sugar for `model.condition(D)`.
+
+**objective** — a scalar function `(model, Dataset) -> ScalarFloat` consumed
+by `gpx.fit`. Objectives are thin: `conjugate_mll` is the evidence view of
+the conditioned posterior, not a second derivation.
+
+**Dataset** — the data container. `n_total` records the full-dataset size
+when the object is a minibatch view (stamped by `get_batch`); the minibatch
+ELBO scale is derived from it, never supplied by hand. Read it through
+`full_size`, which falls back to `n` for a whole dataset — production code
+uses `data.full_size / data.n` and never re-spells the fallback inline.

benchmarks/README.md

@@ -77,6 +77,6 @@ Prints a regression report with significance markers. No branches touched.
 ## Notes on bench-check
 
 `bench-check` runs `asv check`, which builds an env for the commit
-referenced by `asv.conf.json`'s `branches: ["main"]`. It only succeeds
-once `benchmarks/` and `asv-constraints.txt` have landed on `main`. After
-that, it's safe to add to the `all-tests` poe sequence.
+referenced by `asv.conf.json`'s `branches: ["main"]`. It is wired into
+the `all-tests` poe sequence, so it runs on every `uv run poe all-tests`
+alongside `lint`, `docstrings` and `test`.

benchmarks/compile.py

@@ -39,10 +39,10 @@ def setup(self):
         kernel = gpx.kernels.RBF()
         mean = gpx.mean_functions.Zero()
         prior = gpx.gps.Prior(kernel=kernel, mean_function=mean)
-        likelihood = gpx.likelihoods.Gaussian(num_datapoints=n)
+        likelihood = gpx.likelihoods.Gaussian()
         self.posterior = prior * likelihood
         self.q = VariationalGaussian(
-            posterior=self.posterior, inducing_inputs=X[:M_INDUCING]
+            model=self.posterior, inducing_inputs=X[:M_INDUCING]
         )
 
         self.jitted_mll = jax.jit(objectives.conjugate_mll)