V1.0 - #752
Open
thomaspinder wants to merge 106 commits into
Open
Conversation
conjugate_mll had no value-level test anywhere in the suite, yet serves as the oracle for the Kalman MLL and collapsed_elbo. These closed-form pins, computed through an independent jnp.linalg path, give the reference frame its ground truth ahead of the v1.0 conditioning refactor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…tter bug collapsed_elbo(z=X) vs conjugate_mll and whitened-vs-unwhitened predicts at matched parameters now guard the five independent derivations of the conjugate conditioning algebra. The strict xfail documents that at non-default jitter the derivations factorise different matrices — the bug the v1.0 conditioning module removes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
_compare previously swallowed AssertionError with a print, so the harness could never fail. Failures are now collected per-example and raised at the end of test(), making the four golden-value pins a real no-behaviour-change net for the v1.0 refactor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The newly-loud harness exposed pre-existing drift in all four examples: the collapsed/uncollapsed goldens predated the real-data example swap (#696) and the regression/heteroscedastic goldens predated subsequent behaviour fixes (#707/#708/#713 and dependency bumps). The toothless harness never noticed. Re-pinned so the net measures the v1.0 refactor, not history. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The full-dataset size a minibatch ELBO needs now travels on the one object that knows it, as static pytree aux_data, instead of being smuggled through likelihood constructors. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…nd noise_prior Of ~245 occurrences of num_datapoints, only four were real reads: the ELBO minibatch scale (now served by Dataset.n_total) and latent sizing (moves to data-contact time in the JointModel rewrite). No likelihood used the value internally, and nothing validated it — a wrong value silently mis-scaled the ELBO. noise_prior moves to the model layer, where priors live. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
The API now mirrors the maths: prior * likelihood -> JointModel (the joint p(f,y), the trainable object); model.condition(D) — sugar: model | D — returns an immutable Posterior pytree caching the Cholesky factor and representer weights. The predictive, log_marginal_likelihood, loo, and pathwise sample_approx are views of that one factorisation, deleting the eleven independent derivations and the two-owner jitter split (prior.jitter is now the single knob, applied once inside conditioning). - gpjax/conditioning.py: deep module (Posterior, ExactPosterior, LatentPosterior); MO validation moves to condition time; sample_approx refuses multi-output loudly instead of silently broadcasting wrong. - gps.py: Prior (AbstractPrior folded in), ConjugateModel, NonConjugateModel (lazy latent, sized at data contact), HeteroscedasticModel (owns noise_prior — likelihoods are pure conditionals again, killing the likelihoods->gps circular import). Deleted: AbstractPrior, AbstractPosterior, LatentPosterior marker, ChainedPosterior marker, construct_posterior (now construct_model). - objectives: conjugate_mll/conjugate_loocv/log_posterior_density are one-line views of the conditioned posterior. - fit: _prepare_model hook sizes lazily-initialised state from data. - predict(t, D) survives as documented one-line sugar everywhere. - return_covariance_type kwarg renamed to covariance. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Mechanical: ConjugatePosterior->ConjugateModel and friends, construct_posterior->construct_model, return_covariance_type->covariance, num_datapoints/noise_prior constructor ceremony deleted (~240 sites). Semantic: heteroscedastic tests build HeteroscedasticModel directly; non-conjugate tests size the latent via init_latent; docs/index.md quickstart shows condition(); regression example narrates the condition API; StateSpaceConjugatePosterior renamed StateSpaceConjugateModel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…ault jitter The model-side two-owner jitter bug is fixed; the strict xfail narrows to the family-side knob, which unifies in the variational stack PR. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
- reference/gps.md lists the JointModel hierarchy and the conditioning module; state_space.md and linalg.md updated for renamed/new symbols - stale glossary/sharp_bits/classification xrefs renamed - poisson example initialises the lazy non-conjugate latent before MCMC - ADR directory excluded from the docs site (in-repo records for now) - codeautolink match_block warnings suppressed on every path: a matcher limitation on doctest-SKIP blocks, predating this stack — the docs workflow had not run cold since the Sphinx migration, so tonight's PR pushes surfaced it for the first time Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…inery Implements Salimbeni et al. 2018 (arXiv:1803.09151) natural-gradient VI for VariationalGaussian and WhitenedVariationalGaussian. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Adversarial review of the natural-gradient core turned up two correctness defects, one performance defect, three contract mismatches and a set of documentation and test gaps. Fixes, in order of severity: * `_first_valid_trial` leaked float64 into the scan carry. Under x64 the exponent `jnp.arange(K + 1)` is int64, so `backoff ** arange` was a non-weak float64 that promoted a float32 model and made `lax.scan` reject the carry. The trial ladder is now cast to the dtype of Theta_2, so `fit_natgrads` is no longer strictly narrower than `fit`. * The backoff replicated the whole theta -> xi map across all K+1 trials, at a measured 13% of total training wall clock at M=200 -- not the "negligible" cost impl-plan 1.2.5 assumed. Only the admissibility probe is replicated now; the inversion, the X^T X product and the second Cholesky run once, at the accepted step size. Measured overhead at M=200 falls to 4.6%. * `fit_natgrads`' signature rejected everything `_check_natgrad_lr` blessed: `natgrad_lr=1`, `map_jitter=0`, `backoff=1` and a 0-d array all raised under the beartype import hook. Annotations widened, the validator now accepts 0-d arrays and rejects bool, and the entry point (not just the validator) is tested with each. * `_reject_frozen_coordinates` matched coordinates against top-level dataclass fields by identity, so a future nested registration would have passed the guard silently. It now re-walks the tree with the selector as the `is_leaf` predicate and reports the full key path. Its message also pluralises and points at a remedy that exists -- the old one recommended freezing the whole family, which re-raises the same error. * `fit_natgrads` now calls the guard under `safe=True` as impl-plan 1.2.4 prescribes, and forwards `log_rate` to `vscan` instead of documenting a knob that did nothing. Tests: `test_natgrad_backoff_recovers_from_large_step` never exercised the backoff (the k=0 trial was already admissible at natgrad_lr=100), so it now starts from a tight S_0, asserts the un-shrunk step genuinely leaves the cone, and checks the accepted step size. The Cholesky-budget test could not see the vmap width; it is now an exact count parametrised over max_backoff, plus a lowered-IR test asserting the batched Cholesky has leading dimension K+1. Added dtype-preservation tests, and moved the duplicated conjugate oracle into `tests/_reference/conjugate_svgp.py` so the two transcriptions cannot drift. Docs: all seven exported functions gained runnable `Example:` blocks (impl-plan section 6), the map_jitter bias on `history` and the eta -> xi cancellation regime are now documented on the public surface, and two docstrings became raw strings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…_natgrads The ```pycon fences render fine under MkDocs but are invalid RST under Sphinx/napoleon, producing docutils warnings that fail the -W docs-ci gate. Drop the fences; the indented doctest block matches fit()'s style. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
Adds examples/dual_svgp.py, the second of the two natural-gradient tutorial
notebooks, and wires both of them into the documentation.
The notebook derives the dual parameterisation for a reader who has been
through examples/natgrads.py: the additive split eta = eta_0(theta) + lambda,
the EP-style likelihood sites and their tying to inducing space, the two
convention traps (flanked vs un-flanked storage, and the -1/2 on Lambda_2),
and the tied natural-gradient update, which is an affine convex combination on
the stored sites because grad_mu KL == lambda exactly, so the KL is never
differentiated and no theta <-> eta round trip is needed.
Measured in the executed run:
* one rho = 1 full-batch step on a conjugate model with a non-zero mean
function reproduces the Titsias optimum to 1.7e-12 (mean) and 1.2e-13
(covariance), a second step moves nothing, and dual_elbo matches the
collapsed bound up to exactly N * jitter / (2 sigma^2);
* rho is gamma: matched dual and Salimbeni E-steps agree in (m, S) to 3.1e-15
over six full-batch steps at rho in {0.3, 0.8, 1.0}, and two frozen-
hyperparameter fit_natgrads runs overlay to 4.3e-14 over 50 iterations;
* the banana benchmark (N = 2000, M = 50, B = 256, 1000 iterations, the same
make_banana and jr.key(42) as the natural-gradients notebook) with t-SVGP,
the Salimbeni natgrad and Adam alone, per iteration and per wall-clock
second -- both natural-gradient runs reach Adam's 1000-iteration bound at
iteration 109, and the dual iteration is not cheaper at this scale;
* hyperparameter learning: the dual/standard hyper-gradient gap falls from
3.9e1 to 8.5e-15 as the E-step converges; dual dominance is not uniform when
sparse (negative gaps at M = 5 and M = 10) but holds by 1.5e5 nats at Z = X;
and a 40-round VEM loop ends 0.50 nats ahead on dual_elbo with equal
held-out NLPD.
Docs wiring: both notebooks added to the mkdocs.yml Tutorials nav and to
CTA_NOTEBOOKS in docs/scripts/gen_examples.py, plus an adam2021dual entry in
docs/refs.bib.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Corrections to examples/dual_svgp.py, all re-verified against a fresh end-to-end execution: * the intro no longer implies the two hyperparameter gradients agree at theta_t; they agree only at a converged E-step, which is what the notebook's own gradient table measures; * added a notation-reconciliation note bridging the natural-gradients notebook's (theta, eta, lambda) to this one's (eta, mu, lambda), and restated the borrowed identity and H_2 in these letters; * the c(theta) remark now names the site convention it holds under (normalised projected) and fixes the sign apposition; * the bound-slice prose is now asymmetric, as the data are: l collapses on the long-lengthscale side while l-bar barely moves, but both collapse together on the short side, where the sparse approximation itself has failed. Added edge diagnostics to back it; * the banana benchmark now states which run ends ahead and bounds what that comparison can mean; * the VEM panel plots the round-by-round bound lead rather than two indistinguishable traces. That exposed a false claim: the dual M-step is behind for the first seven rounds, crosses at round 8 and holds a sub-nat lead thereafter. Prose corrected and the crossing printed; * order-of-magnitude claims restated from the printed values (Titsias agreement, the jitter residual now printed to twelve digits, the banana condition-number ratio, the M = 20 crossing-point noise); * the dominance row now names conjugacy as well as Z = X; * the roadmap names the two sections it had omitted; * the banana-copy rationale no longer overstates cross-notebook comparability, and the wall-clock explanation leads with the O(M^3)-vs-O(BM^2) argument rather than an evaluation XLA folds away; * ruff-format clean, no code line over 88 characters. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
…tionVariationalGaussian Both classes were parameterisation-only: they stored the natural or expectation coordinates of q(u) but shipped no way to take a natural-gradient step in them, so they bought nothing over the standard families. Natural-gradient geometry belongs to the optimiser, not the family. The Fisher matrix is exactly the Jacobian dn/dt, so the natural gradient with respect to the natural parameters equals the ordinary gradient with respect to the expectation parameters, in any parameterisation. fit_natgrads (PR#1, gpjax/natural_gradients.py) therefore computes the transforms on the fly and operates directly on VariationalGaussian and WhitenedVariationalGaussian, which store constraint-respecting coordinates. Users of the removed classes should switch to VariationalGaussian with gpjax.fit_natgrads. Also drops the now-dead _psd helper and the cholesky_factor import, whose only call sites lived inside the deleted classes, and the "natural" and "expectation" arms of the VariationalParametrisationSuite ASV benchmark. BREAKING CHANGE: NaturalVariationalGaussian and ExpectationVariationalGaussian are removed from gpjax.variational_families. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Implements the dual/site parameterisation of sparse variational GPs from Adam, Chang, Khan and Solin (2021), "Dual Parameterization of Sparse Variational Gaussian Processes", NeurIPS 2021 (arXiv:2111.03412). `DualVariationalGaussian` stores an unnormalised Gaussian site on the *centred* inducing outputs -- `dual_vector` is the site's first natural parameter and `dual_matrix` its precision, both `Real` and both defaulting to zero, so q(u) = p(u) at initialisation. Neither carries a constraining bijection: PSD-ness of the site precision comes from the convex-combination structure of the natural-gradient update, and a bijection would destroy that affine step. Moments, marginals, prior KL and predictions all route through the working matrix R = Kzz + Kzz L2 Kzz, which dominates Kzz and is therefore always factorisable even when the site precision is rank deficient; exactly two Cholesky factorisations are taken per call and nothing is inverted. The centred convention is the correction to the reference implementation's mean-function bug, which shifts by `predict_f(Z)` and so is wrong for any non-zero mean function. `test_dual_natgrad_handles_non_zero_mean_function` pins the correct behaviour and asserts that the uncentred variant misses. `marginals` adds the family's jitter to every marginal variance. This is load-bearing, not cosmetic: `VariationalGaussian.predict` runs `add_jitter` on its output covariance, so the per-point marginals `elbo` sees carry the same offset, and without it `dual_elbo` would miss `elbo` at matched moments by N*eps/(2 sigma^2). `dual_elbo` is the same functional as `elbo` but evaluated as a function of the sites and the hyperparameters. Its value matches `elbo` at the implied moments (measured 5.7e-14 absolute, 2.8e-16 relative at random PSD sites) while its hyperparameter gradient differs, because q moves with theta through Kzz while the sites stay frozen. Kzz is deliberately not detached and no moments are cached on the module; a cached implementation would pass every value assertion and fail only `test_dual_elbo_hyper_gradients_differ_away_ from_optimum`. `natural_gradient_step` and `variational_coordinates` gain dual registrations. Because grad_mu KL = lambda exactly, the KL is never differentiated and the step is a convex combination towards a closed-form target built from one `jax.grad` of `expected_log_likelihood` (Bonnet and Price), with N/B scaling, a trace-safe floor on beta and a symmetrise. That makes it the Salimbeni step at gamma = rho: matched initialisations agree to 1.4e-15 in (m, S) over six Bernoulli steps at every rate tested. One rho = 1 full-batch conjugate step lands on the Titsias optimum to 4.9e-15. `fit_natgrads` rejects a numeric step size above one on this family, since the update is a convex combination; schedules cannot be checked statically and are documented as the caller's responsibility. Plain `fit` on the family also works and is documented as gradient descent in the dual coordinates. The `VariationalParametrisationSuite` benchmark gains a `dual` arm, in both `benchmarks/objectives.py` and the hardcoded parametrize list in `tests/test_benchmarks_smoke.py`. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Adds examples/natgrads.py, the first of two tutorial notebooks for the natural-gradient stack. It derives the exponential-family view of q(u), shows that the Fisher information is the Jacobian d(eta)/d(theta) (checked numerically to 2.9e-15), reads the update as mirror descent, and then runs two demos: * a conjugate 1D regression where one gamma=1 full-batch natural-gradient step recovers the Titsias optimum to 1.2e-13 while Adam on the same problem is still 4.6 nats short after 2000 iterations; * a mini-batched 2D banana Bernoulli benchmark (N=2000, M=50, B=256, 1000 iterations) comparing natural gradients + Adam against Adam alone, per iteration and per wall-clock second. Closes with the negative-definite cone result, a gamma sweep reproducing its boundary, and a demonstration of the step-size backoff. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
`mkdocs build` aborted with `IndexError: string index out of range` while
rendering `_examples/dual_svgp.md`. The trigger is a markdown-katex parser
bug: `iter_inline_katex` reads `line[end + 1]` without a bounds check, so any
line whose final characters are a backtick code span immediately preceded by
`$` crashes the build. The prose read
... that is $-$`sparsity_gap`
above, and ...
where the closing `$` of `$-$` abuts the code span and the span ends the line.
Reword to "the negated `sparsity_gap` computed above", which removes the
`$`-adjacent code span entirely rather than relying on a particular line wrap.
The meaning is unchanged, and the sentence still says c(theta) is minus the
Titsias trace term. A scan of all generated `docs/_examples/*.md` confirms this
was the only occurrence of the pattern in the docs tree.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
CHANGELOG: add the missing "### Added" entry for fit_natgrads and gpjax.natural_gradients. PR#1 shipped both without a changelog entry, so the Removed entry added here forward-referenced an API the changelog never announced. Bringing it forward from PR#3 keeps any release cut mid-stack self-consistent; PR#3 appends the dual entries to the same section. CHANGELOG: correct the justification prose. "the natural gradient with respect to theta equals the ordinary gradient with respect to eta, in any parameterisation" is false as literally written -- for a reparameterisation xi with J = dtheta/dxi, the natural gradient in xi is J^-1 grad_eta L, not grad_eta L. The identity is specific to the natural/expectation pair of an exponential family. Reworded to state that pairing and the point it supports: either coordinate system is recoverable on the fly, so no dedicated class is needed. benchmarks: drop diff-relative wording from the VariationalParametrisationSuite docstring. "surviving" only means something to someone reading this commit's diff, and "Both" is a count PR#3 invalidates when it re-adds the dual arm. The module docstring keeps its explicit "(standard, whitened)" list, which PR#3 must extend regardless. tests: split the _psd guard out of test_removed_families_are_gone. The _psd arm asserted `"_psd" not in __all__`, vacuous for a helper that was never exported, under a failure message about superseded parameterisations. It is now its own test with a docstring saying what it actually guards. CLAUDE.md: "Three optimisers" -> four. fit_natgrads landed in gpjax/fit.py in PR#1; the sentence has been stale since. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Factor R in the Kzz basis. Forming R = Kzz + Kzz Lambda_2 Kzz explicitly carries a rounding error of order ||Kzz||^2 ||Lambda_2|| eps, which for a large-variance kernel (or in float32) exceeds lambda_min(R) ~= jitter, so chol(R) returned NaN and poisoned every later fit_natgrads iterate -- measured at RBF variance 1e4, M=80, default jitter, float64, where the matched VariationalGaussian run stayed finite. R = Lk (I + Lk^T Lambda_2 Lk) Lk^T is factorised instead, giving a lower-triangular Lr = Lk chol(I + G) whose inner matrix has lambda_min >= 1 - O(||G|| eps). Same two Choleskys per call, plus one M x M product. The "chol(R) never fails" and "no Cholesky a backoff could rescue" claims are softened to match. Also: split _gram_and_root off _working_matrices so the dual natgrad step stops discarding a chol(R); take tr(R^-1 Kzz) as ||Lr^-1 Lk||_F^2 rather than a full cho_solve; bound-check an optax schedule against rho <= 1 over the whole num_iters horizon for the dual family, which previously returned a silent all-NaN history; hoist the _fmt_Kzt_Ktt/_fmt_inducing_inputs hooks to AbstractVariationalGaussian and keep one typed _symmetrise; convert the dual family's numpydoc sections to the Google style the file and mkdocs use, and document marginals' inputs argument. Doc corrections, all measured: elbo on a DualVariationalGaussian returns the same value and the same gradients as dual_elbo (bit-identical value, 1.7e-14 on gradients), so the CHANGELOG's gradient claim now names the matched VariationalGaussian as the comparison; and vmap does not repeat the unbatched factorisations per datum, so neither elbo nor the benchmark arm pays 2N Choleskys -- eager counts are 4 (dual), 2 (standard), 1 (whitened), and the compiled dual_elbo and dual step are 2 potrf each. Tests: regression for chol(R) at variance 1e4/M=80 and 1e3/M=50 with the default jitter, an Lr Lr^T = R reconstruction check, the schedule guard, a half-batch arm on the dual/elbo equivalence plus a direct pin on the N/B factor, and the triplicated dual fixtures moved to tests/_dual_helpers.py. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
Correct the mini-batch ramp argument (the target is q-dependent outside conjugacy, so gamma=1 lands on a moving fixed-point target, not the mini-batch optimum), replace the unmeasured calibration claim after the banana contours with the metrics the cell actually prints, and attribute the cone sweep's gamma=2 failure to the over-confident S_0 rather than to gamma=2 itself. Smaller corrections: the Fisher solve is O((M + M(M+1)/2)^3) in the vec_s coordinates, not O((M + M^2)^3); the one-step demo agrees to ~1e-13, not fourteen decimal places; the sparse/exact predictive deviation is quantified and located outside the data range; the Adam ELBO-gap description now matches the shape of the log-log panel; the roadmap says "exact variational optimum" where the notebook later reserves "exact posterior" for the non-sparse GP; K=100 is explained against the paper's dataset-dependent K. Code: derive the ELBO figure's y-limits from the smoothed histories so neither curve is clipped, guard the crossing report against never reaching the target, draw the held-out points in the categorical palette with per-class markers instead of the contour colourmap, use banana_data in the split report, and run ruff format (the pinned make_banana block is byte-identical afterwards). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
The dual/Salimbeni E-step divergence on the banana demo was attributed to floating-point conditioning. It is not: `inv_probit` clips its output into [1e-3, 1-1e-3], so the computed Bernoulli log-likelihood is not log-concave in the tails (positive second derivative for f < -2.44). A confidently mislabelled point then yields beta_i < 0, the dual branch's beta_floor clips it, and the two branches diverge. With the clip disabled the same six steps agree to 6.2e-13 instead of 5.1e-3. - Re-attribute the mechanism in the dual notebook and add a diagnostic cell that measures it, and qualify the "identical iterates" claim wherever it is stated (natural_gradients.py, fit.py, CHANGELOG, both notebooks): it holds provided the computed beta stays non-negative. - Note in the natgrads notebook that the cone discussion assumes a log-concave *computed* likelihood, which the clipped probit violates in the far tails. - Reword the Salimbeni registration docstrings: dispatch covers GraphVariationalGaussian, but the standard elbo path is broken upstream (MatrixLinearOperator dimensionality error out of gram), on main too. - Extend _check_natgrad_schedule to reject non-positive rates for every family, not just the dual upper bound. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHxYz2P7JCSqpax5RmDwWD
The Sphinx docs arrived with v1.0, after this branch was cut, so the deletion of NaturalVariationalGaussian and ExpectationVariationalGaussian now has to reach three doc files the original commits could not know about: drop the two classes from the variational-families autosummary page, and repoint the glossary's "natural parameters" entry from the removed classes to fit_natgrads on the surviving families. fit_natgrads is added to the fit reference page so that glossary link resolves — an omission from PR#1, which shipped the function without a reference entry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
…e v1.0 API The v1.0 likelihoods are pure conditionals, so the minibatch ELBO scale in dual_elbo and the dual natural-gradient step now derives from Dataset.n_total instead of likelihood.num_datapoints, the constructor annotation follows the AbstractPosterior -> JointModel split, and the docstring examples plus the dual test fixtures drop num_datapoints. The minibatch tests stamp n_total onto their batch views to preserve the N/B factor they assert. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019d7TF7oQt2Du4EQ74bMBQB
SparsePosterior gains sample_approx, the Wilson et al. (2020, section 5.2)
decoupled sampler for sparse variational posteriors: a prior draw from
Fourier features of the kernel, corrected by K_zt Kzz^{-1}(u - f_prior(z))
where u ~ q(u) is a draw at the inducing inputs. Reuses the posterior's
cached cholesky_kzz, variational_mean and variational_root -- no extra
factorisation -- and handles both the unwhitened and whitened
parameterisations via the existing `whitened` flag, de-whitening the u
draw through the same cached Cholesky factor before the correction.
Mirrors ExactPosterior.sample_approx's structure and its multi-output
refusal (ValueError, not a silent broadcast) for consistency across the
module. GraphVariationalGaussian is left unsupported by construction: RFF
requires a StationaryKernel and GraphKernel is not one, so it fails with
the existing, clear TypeError rather than a new bespoke guard.
Tests (tests/test_conditioning.py) cover VariationalGaussian and
WhitenedVariationalGaussian: Monte Carlo mean/covariance agreement with
the analytic predictive, determinism in the PRNG key, and jit/grad/vmap
cleanliness.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
… posterior Implements the acceptance criteria from #651: `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and `StateSpacePosterior.__call__` now support `covariance="dense"`, matching the dense `ConjugateModel` predictive (mean and full covariance) to ~1e-8 for Matern-1/2, 3/2, 5/2. - `rts_smoother` gains an opt-in `return_gains=True` that exposes its already-computed per-step smoother gains (no extra numerical work, just not discarding them). Existing 2-tuple callers are unaffected. - New `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains those gains into the M x M joint covariance across test points, following the RTS smoother cross-covariance recursion (Sarkka & Solin 2019 SS12.2): Cov(x_i, x_j | y) = G_i...G_{j-1} P_j^smooth for i < j. It never inverts a gain product (ill-conditioned for widely separated points, since gains shrink with lag) and never forms an N x N gram over the training set -- cost is O(N d^3) for the filter/smoother pass plus O(M^2 d^3) for the cross-covariance chaining, both linear in N. - `StateSpacePrior.predict` (no conditioning data) returns the kernel's own dense gram for `covariance="dense"`, since the SDE is an exact representation of the kernel -- no Kalman machinery needed. Design decision (the issue underspecifies this): `StateSpacePosterior.filtered` / `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep raising NotImplementedError for `covariance="dense"`. Each filtered test point 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 -- the issue's acceptance criteria only compares against the dense conjugate predictive, which is the smoothed quantity, and an existing repo test (test_state_space_posterior_predict_filter_dense_raises) already pinned the filtered-raises behaviour. Extending the filtered path is left to a future issue if there's demand. Tests: mean+covariance equivalence to the dense ConjugateModel (parametrized over Matern12/32/52 and jitter), diagonal/dense consistency, caller-order preservation, the M=1 degenerate case, a joint-sampling smoke test, jit and grad cleanliness, a machine-precision numpy-oracle check on the exposed smoother gains, and a larger-N robustness smoke test. Existing tests that pinned the old "dense raises" behaviour for the now-supported paths are updated to positive equivalence/plumbing checks instead. Note for reviewers: gpjax/state_space/inference.py::rts_smoother is implemented against its current covariance-form recursion, not a square-root rework (that's issue #668, tracked separately). If #668 lands a true square-root smoother, this cross-covariance recursion is worth revisiting -- a square-root form may simplify the derivation -- but nothing here blocks on it. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Closes #669. The to_sde singledispatch handler for ProductKernel previously rejected every product unconditionally, citing Kronecker state-dimension blowup. That blanket rejection was overbroad: the canonical quasi-periodic kernel TruncatedPeriodic x Matern (Solin & Sarkka 2014 sec. 3) has state dimension (2K+1)*d, not the product of arbitrary factor dimensions, because the periodic factor's transition is an exact rotation. Adds ProductSDE, a Kronecker-sum composition of exactly two factor SDEs (periodic_factor, matern_factor). F is the Kronecker sum F_p @ I + I @ F_m; L, Qc, H, and the stationary covariance sqrt are built from the two factors' own fields and satisfy the continuous Lyapunov identity exactly (verified by test_product_sde_lyapunov_identity). Design fork worth flagging: discretise() does NOT use a generic eigendecomposition-based square root of P_inf - A P_inf A^T (the approach SumSDE's Matern components use via _psd_sqrt). The periodic factor's per-harmonic eigenvalues repeat in pairs (cos/sin components share variance), which makes the combined stationary covariance's spectrum degenerate and jnp.linalg.eigh's reverse-mode gradient singular there (verified this empirically: a naive eigh-based sqrt at the product level produced NaN gradients through the Matern lengthscale). Instead, discretise() exploits two exact identities: A(dt) = A_periodic(dt) (x) A_matern(dt) (Kronecker-sum exponentials commute), and, because TruncatedPeriodicSDE is provably lossless (A_p(t) P_inf,p A_p(t)^T = P_inf,p for all t -- exactly what test_truncated_periodic_L_Q_is_zero already asserts), the discrete noise factorises as Q(dt) = P_inf,p (x) Q_matern(dt), reusing the Matern factor's own already-gradient-safe discretise() instead of any new eigendecomposition. This makes ProductSDE specific to "one lossless factor times one diffusive factor" rather than a fully general two-factor Kronecker composition; to_sde's dispatch enforces the (periodic_factor, matern_factor) argument order this relies on, and the genuinely-unsupported paths (Matern x Matern, three-factor products, non-truncated Periodic x anything) still raise NotImplementedError with the Kronecker-blowup message. Updates the existing blanket-rejection test to keep asserting rejection for genuinely unsupported products, and adds coverage for: state dimension, the Lyapunov identity, A(dt) matching both the per-factor Kronecker product and scipy's matrix exponential of the full drift, zero-dt behaviour, jit/vmap compatibility, gradient finiteness, and kernel values matching the dense product kernel. An end-to-end test confirms StateSpaceConjugatePosterior.predict on TruncatedPeriodic() * {Matern12,Matern32,Matern52}() matches a dense gpx.gps.Prior posterior with the same product kernel to 1e-4 tolerance. docs/examples/state_space_gps.py's Mauna Loa example still sums a seasonal TruncatedPeriodic term with a Matern52 trend and Matern32 wiggle rather than using the new product; swapping the seasonal term to a real TruncatedPeriodic * Matern52 product is now a viable follow-up but is out of scope here per the issue. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…perators GaussianDistribution.kl_divergence calls lx.linear_solve(..., solver= lx.Triangular()), and Triangular.init queries lx.has_unit_diagonal on the operator. BlockDiag and Kronecker had the other 7 Lineax singledispatch predicates registered (is_symmetric, is_diagonal, is_tridiagonal, is_lower_triangular, is_upper_triangular, is_positive_semidefinite, is_negative_semidefinite) but not this one, so any KL divergence between covariances backed by either operator raised NotImplementedError (#709). Both are registered to return False, unconditionally: neither operator guarantees a unit diagonal from its constituent blocks/factors in general, matching the same conservative default Lineax uses for its own DiagonalLinearOperator and TridiagonalLinearOperator. Grepped the codebase for other lx.is_*/lx.has_* predicate queries and other lx.linear_solve call sites reachable with BlockDiag/Kronecker operators; distributions.py:287 is the only such call site, so no other gaps exist. Extends tests/test_distributions_kl.py with a KL round-trip covering both operators: numerical equivalence against an independently-computed dense reference, the public kl_divergence() method, self-KL-is-zero, and jit-compatibility. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Issue #668: rts_smoother's docstring claimed "square-root" but the backward step materialised P = L @ L.T (covariance form), did a Joseph-style covariance update, then re-rooted via _psd_sqrt (an eigh call) every step -- an O(d^3) eigendecomposition per step that defeats the conditioning benefit square-root filtering is supposed to give, and produces non-triangular V·Λ^½ factors despite the docstring's claim. Replace the backward recursion with a QR pre-array square-root smoother (Park & Kailath 1996). Each backward step treats the RTS gain computation as a virtual matrix-valued Kalman measurement update (observation matrix A_{i+1}, observation noise Q_{i+1}) using the same generalised Potter/Bierman QR pre-array as the existing _sqrt_update, then closes the additive covariance term with a second QR combine. Neither step ever forms L @ L.T or calls eigh; both blocks reuse the same "stack transposed factors, QR, transpose R back" pattern, factored out as a new _qr_sqrt_sum helper that _sqrt_predict is also refactored to use (no behaviour change there). Design note: Ls_predicted from _sqrt_filter_forward's output is no longer consumed by rts_smoother -- the QR pre-array reconstructs the equivalent quantity (R11) itself from a fresh sde.discretise(time_step_next) call, exactly mirroring how the prior implementation already recomputed transition_matrix_next rather than threading it through from the forward pass. The forward_outputs tuple shape is left unchanged for API compatibility with prediction.py's two call sites. Verification: existing correctness tests (dense-GP-posterior match, machine-precision NumPy reference, near-noiseless finiteness) pass unchanged at the same tolerances, confirming this is purely an implementation-form change. Added test_rts_smoother_factors_are_lower_triangular, using a state_dim=2 Matern-3/2 SDE (a state_dim=1 SDE would make triangularity trivially true), which the previous eigh-based V·Λ^½ factors could not have passed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…tation AbstractKernelComputation.gram() called self.cross_covariance(kernel, x, x) directly instead of self._gram(kernel, x), leaving _gram (and any subclass override of it) dead code. This silently broke BasisFunctionComputation's _gram override, which computes the RFF feature matrix once: every RFF.gram(x) call was routing through cross_covariance's _cross_covariance override instead, which independently computes the feature matrix twice (once as z1, once as z2) even though x is the same array both times. gram() now calls self._gram(kernel, x), with the base class's _gram implemented as cross_covariance(kernel, x, x) -- identical behaviour to before for every kernel that doesn't override _gram (RBF, Matern, etc. via DenseKernelComputation). DiagonalKernelComputation and ConstantDiagonalKernelComputation override gram() itself, not _gram, so they are unaffected. RFF now correctly uses its single-pass _gram override, computing the feature matrix once instead of twice. Adds a monkeypatched call-count test that fails before this fix (2 calls) and passes after (1 call), plus numerical-equivalence tests confirming gram() output is unchanged for RBF, Matern12/32/52, and RFF kernels. Fixes #679. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…init workaround, fix White's phantom lengthscale Scoped fix for issue #695 per plans/2026-07-26-issue-695-abstract-final.md (the "measured middle path", not a wholesale abstract/final rewrite). 1. Give StationaryKernel/AbstractKernel fields real defaults matching their __init__ defaults (lengthscale/variance via eqx.field(default_factory=...), compute_engine via DenseKernelComputation). Equinox inherits hand-written __init__ methods at runtime, but pyright's @dataclass_transform synthesises a fresh __init__ per subclass from declared fields whenever a class (e.g. RBF, Matern12/32/52) doesn't define its own -- so undefaulted fields made the canonical `RBF()` a pyright error while `RBF(name="xyz")` type-checked cleanly (TypeError at runtime). `name` is now a ClassVar on every concrete stationary kernel, matching runtime (it was never a real __init__ param) and removing it from pyright's synthesised signature. Note: ClassVar must be imported from stdlib `typing`, not `beartype.typing` -- despite `beartype.typing.ClassVar is typing.ClassVar` at runtime, pyright's dataclass field-exclusion does not recognise the re-exported symbol and keeps such fields in the synthesised __init__. Verified against pyright 1.1.411. 2. Removed `_compute_base_init` (kernels/base.py). Its docstring claimed equinox modules are frozen once any parent __init__ returns, which no longer holds under equinox 0.13.8 (verified directly: a child class can set its own fields after `super().__init__()` returns). StationaryKernel now calls plain `super().__init__(...)`. 3. White hardcoded lengthscale=1.0 into StationaryKernel.__init__ even though White.__call__ (and its inherited spectral_density, which raises before touching lengthscale) never reads it, so every White carried a real, trainable PositiveReal leaf with zero gradient that polluted optimiser/MCMC state. White now overrides `lengthscale` as a ClassVar (`None`) rather than inheriting it as a dataclass field, and calls `AbstractKernel.__init__` directly instead of `StationaryKernel.__init__`, so lengthscale is absent from White's pytree entirely (1 leaf instead of 2 -- verified via jax.tree_util.tree_flatten). The Zero mean-function fix from the same research note was already in place (gpjax/mean_functions.py) and is unaffected by this change; reverified via the note's repro script. Residual, out of scope: SumKernel/ProductKernel (kernels/base.py) still synthesise a `kernels: tuple` parameter under pyright that rejects a `list[AbstractKernel]` argument (`SumKernel(kernels=[...])`). This was already broken pre-fix (masked by RBF()'s own missing-argument errors) and sits one level deeper -- fixing it needs a converter on the `kernels` field, which is a different kind of change to the one scoped here. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…683) Both functions implemented the identical power-sum/sign/mask/fori_loop Newton-Girard recursion, differing only in the trailing broadcast shape of the per-dimension values (scalar in oak.py vs. (N, N) matrix in sobol.py). Generalised gpjax/kernels/additive/oak.py::_newton_girard to accept an arbitrary trailing shape via broadcasting (exponents/signs/mask reshaped against z.ndim, elem_sym shaped from z.shape[1:]), and made sobol.py import and call it directly, deleting _newton_girard_matrices. Pure refactor: replaced the scalar-only jnp.dot reduction with jnp.sum(product, axis=0), which computes the identical sum and matches what the matrix variant already did, so the scalar call site's output is numerically unchanged. Verified via existing OAK/Sobol tests (unchanged, all passing) plus a new test asserting the matrix-shaped path reduces to the same per-element scalar recursion. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
The soak-period precondition from #637's rollout plan (benchmarks/ and asv-constraints.txt landing on main) has been satisfied for months, so per plans/2026-05-01-continuous-benchmarking-impl.md Task 12, bench-check now runs as part of `uv run poe all-tests`, between docstrings and test. Updates benchmarks/README.md's "Notes on bench-check" section, which previously explained why the task wasn't yet wired in, to describe the current wiring instead. Adds tests/test_poe_tasks.py asserting the all-tests sequence and bench-check task definition via tomllib, following the existing pyproject.toml-parsing pattern in tests/test_dependencies.py. Closes #638. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
…kelihood/covariance API The fix commit was cherry-picked from a worktree that was accidentally based on main instead of v1.0; its own new test still used the pre-v1.0 Gaussian(num_datapoints=...) constructor and the return_covariance_type= kwarg, both removed by the v1.0 rename sweep. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Closes #515. Adds a way to regularise gradient-based MLE fitting with hyperparameter priors (GPyTorch-style MAP), e.g. strongly preferring large lengthscale/noise over overfitting-prone small ones in high dimensions. Design: `with_log_prior(objective, log_prior)` is a plain combinator over the existing `Objective = Callable[[Model, Dataset], ScalarFloat]` protocol -- it returns a new objective computing `objective(model, data) + log_prior(model)`, where `log_prior` is a user-supplied function over the model pytree. This composes directly with `fit`/`fit_scipy`/`fit_lbfgs` with no changes to those functions, and conjugate_mll/other objectives are untouched when the feature is unused. This deliberately does NOT reintroduce the pre-v1 `Parameter(..., prior=...)` field removed in 98a7feb (#621). That mechanism attached a prior to every Parameter class, which (a) tangled the constrained/ unconstrained bijection with an ambiguous "is this prior for gradient descent or for NumPyro?" scope (the numpyro-priors.md design doc that preceded the removal explicitly called this out and worked around it with a separate `numpyro_properties` namespace), and (b) duplicated the fully Bayesian path, which now samples hyperparameters directly via `numpyro.sample` and feeds raw arrays into GPJax constructors (see `tests/test_numpyro_extras.py`). A standalone objective combinator avoids both problems: no field on Parameter, no coupling to NumPyro, and the fully-Bayesian path is untouched. Known simplification, documented in the docstring: `log_prior` is evaluated on the constrained (unwrapped) parameter values -- the same values `objective` receives -- without the change-of-variables Jacobian for the unconstrained space the optimiser actually moves in. For the strongly regularising priors this feature targets, that's an acceptable simplification rather than a literal Bayesian MAP. Tests added to tests/test_objectives.py: - test_with_log_prior_matches_base_objective_when_prior_is_zero: a zero log-prior leaves both value and gradient identical to the unregularised objective. - test_with_log_prior_map_regularised_fit_prefers_prior_consistent_lengthscale: end-to-end MAP demonstration -- an unregularised fit_scipy fit on a high-frequency, low-noise signal collapses the RBF lengthscale to a tiny, overfitting value, while wrapping the same objective with a LogNormal(log(3), 0.15) prior on the lengthscale converges near the prior mean instead. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
Adds a heavy-tailed Student's t observation likelihood (Jylanki, Vanhatalo & Vehtari 2011) following the Bernoulli/Poisson pattern: GHQuadratureIntegrator(20) for the expected log-likelihood, since the Student's t is not conjugate to a Gaussian latent. degrees_of_freedom and scale are both PositiveReal-wrapped. Ported from a fix implemented against a stale pre-v1.0 checkout (the original used the removed Gaussian-style num_datapoints constructor argument); adapted here to v1.0's pure-conditional likelihood contract and NonConjugateModel.init_latent() lazy-latent pattern. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
|
📖 Docs preview: https://pr-752--endearing-crepe-c2d5fe.netlify.app Smoke render — the expensive notebooks run with reduced budgets, so |
Fix bug when unwrapping paramax parameters
# Conflicts: # gpjax/state_space/inference.py # tests/test_state_space/test_prediction.py
…ovariance feat(state-space): dense joint predictive covariance
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Checklist
uv run poe formatbefore committing.Description
Please describe your changes here. If this fixes a bug, please link to the issue, if possible.
Issue Number: N/A