docs(examples): dual/t-SVGP tutorial notebook + docs wiring [stack 7/10]

.github/workflows/security-analysis.yml

@@ -3,8 +3,14 @@ name: "Security Analysis"
 on:
   push:
     branches: [ "main" ]
+  # Deliberately unfiltered: the two jobs below are required contexts on `main`,
+  # and a `branches: [main]` filter here meant they only ever reported on PRs
+  # that already targeted `main`. Stacked PRs — which target their predecessor,
+  # not `main` — produced 16 checks instead of 18, so the whole of a stack got
+  # reviewed and approved without ever being secret- or dependency-scanned; the
+  # gate first fired at retarget time, after review. Scan the code, not the
+  # target branch.
   pull_request:
-    branches: [ "main" ]
   schedule:
     - cron: '30 4 * * 1'  # Weekly on Mondays at 4:30 AM UTC
 
@@ -75,4 +81,18 @@ jobs:
       uses: trufflesecurity/trufflehog@main
       with:
         path: ./
-        extra_args: --debug --only-verified
+        # `--exclude-detectors=lob`: TruffleHog's Lob key pattern is
+        # `(live|test)_` followed by 35 more characters, which matches any
+        # pytest function name that happens to be exactly 40 characters long —
+        # and its verifier then reports those as *verified*, so `--only-verified`
+        # does not filter them out. Four real examples in this repo:
+        # test_fit_natgrads_accepts_optax_schedule,
+        # test_dual_natgrad_matches_salimbeni_step,
+        # test_dual_working_matrices_reconstruct_r,
+        # test_prior_kl_matches_textbook_reference.
+        # Reproduced standalone: a fresh repo containing only one of those names
+        # yields `Lob verified=true`. Since `Secrets Detection` is a required
+        # context on `main`, leaving this on means an arbitrary, length-dependent
+        # subset of PRs can never merge. GPJax has no Lob (direct-mail API)
+        # integration, so excluding the detector costs no real coverage.
+        extra_args: --debug --only-verified --exclude-detectors=lob

CHANGELOG.md

@@ -8,6 +8,62 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **`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.
+
+### 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
 
 - **`Zero` mean function is trainable and drifts away from zero.** Fitting a

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,62 @@
+# 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").
+Signatures follow the maths: models take data; already-fit variational
+families will condition without arguments (post-natgrads stack PR).
+
+**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.
+
+**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.

benchmarks/compile.py

@@ -39,7 +39,7 @@ 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]

benchmarks/objectives.py

@@ -9,9 +9,9 @@
 both M (inducing count) and batch_size; VfeElboSuite (collapsed analytic
 ELBO) varies M only — the collapsed objective requires the full dataset.
 
-VariationalParametrisationSuite compares the four variational Gaussian
-parameterisations (standard, whitened, natural, expectation) at one
-fixed (n, M) so users can see the per-step cost of the choice.
+VariationalParametrisationSuite compares the three variational Gaussian
+parameterisations (standard, whitened, dual) at one fixed (n, M) so users
+can see the per-step cost of the choice.
 
 HeteroscedasticElboSuite and OilmmPredictSuite are independent — they
 do not participate in the alignment because their model structure
@@ -26,9 +26,8 @@
 from gpjax.likelihoods import HeteroscedasticGaussian
 from gpjax.variational_families import (
     CollapsedVariationalGaussian,
-    ExpectationVariationalGaussian,
+    DualVariationalGaussian,
     HeteroscedasticVariationalFamily,
-    NaturalVariationalGaussian,
     VariationalGaussian,
     WhitenedVariationalGaussian,
 )
@@ -50,7 +49,7 @@ def _conjugate_posterior(n: int):
     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()
     return prior * likelihood, data
 
 
@@ -117,17 +116,25 @@ def time_collapsed_elbo(self, n, m):
 _VARIATIONAL_FAMILIES = {
     "standard": VariationalGaussian,
     "whitened": WhitenedVariationalGaussian,
-    "natural": NaturalVariationalGaussian,
-    "expectation": ExpectationVariationalGaussian,
+    "dual": DualVariationalGaussian,
 }
 
 
 class VariationalParametrisationSuite:
-    """Per-step ELBO cost across the four variational Gaussian families.
-
-    All four parameterise the same q(u); the differences are in how the
-    KL term and predictive moments are computed. Holding (n, M) fixed
-    isolates the parameterisation cost.
+    """Per-step ELBO cost across the variational Gaussian parameterisations.
+
+    All three parameterise the same q(u); the differences are in how the
+    KL term and predictive moments are computed. The dual (t-SVGP) family
+    stores sites rather than moments and recovers q(u) through
+    R = Kzz + Kzz Lambda_2 Kzz, so every predict and every KL factorises
+    both Kzz and R, against one factorisation for the standard family and
+    none in the whitened KL. Holding (n, M) fixed isolates that cost.
+
+    The arm deliberately times the generic ``elbo`` on all three families
+    rather than ``dual_elbo`` on the dual one: the point of the comparison
+    is the parameterisation, so the objective has to be held fixed. The
+    dual family's own fast path is ``dual_elbo``, which replaces the
+    per-point ``predict`` with one batched ``marginals`` call.
     """
 
     params = (list(_VARIATIONAL_FAMILIES),)
@@ -157,7 +164,7 @@ def setup(self, n):
         noise_prior = gpx.gps.Prior(
             kernel=noise_kernel, mean_function=gpx.mean_functions.Zero()
         )
-        likelihood = HeteroscedasticGaussian(num_datapoints=n, noise_prior=noise_prior)
+        likelihood = HeteroscedasticGaussian(noise_prior=noise_prior)
         posterior = signal_prior * likelihood
         Z = data.X[:M_INDUCING]
         self.q = HeteroscedasticVariationalFamily(

benchmarks/state_space.py

@@ -9,7 +9,7 @@
 
 import gpjax as gpx
 from gpjax.state_space import StateSpacePrior, state_space_mll
-from gpjax.state_space.gps import StateSpaceConjugatePosterior
+from gpjax.state_space.gps import StateSpaceConjugateModel
 import jax.numpy as jnp
 import jax.random as jr
 
@@ -35,10 +35,8 @@ def setup(self, n):
             mean_function=gpx.mean_functions.Zero(),
             kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
         )
-        likelihood = gpx.likelihoods.Gaussian(num_datapoints=n, obs_stddev=0.1)
-        self.posterior = StateSpaceConjugatePosterior(
-            prior=prior, likelihood=likelihood
-        )
+        likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1)
+        self.posterior = StateSpaceConjugateModel(prior=prior, likelihood=likelihood)
         self.data = _temporal_dataset(n)
         realise(state_space_mll(self.posterior, self.data))