Add functionality for MAP estimation of parameters - #774
Open
Thomas-Christie wants to merge 112 commits into
Open
Thomas-Christie wants to merge 112 commits into
Thomas-Christie wants to merge 112 commits into
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 (QuantClimate#696) and the regression/heteroscedastic goldens predated subsequent behaviour fixes (QuantClimate#707/QuantClimate#708/QuantClimate#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
github-actions
Bot
added
documentation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| if: github.event_name == 'pull_request' | ||
| uses: codelytv/pr-size-labeler@v1 | ||
| if: github.event_name == 'pull_request_target' | ||
| uses: codelytv/pr-size-labeler@095a41fca88b8764fd9e008ad269bcdb82bb38b9 # v1 |
This branch has not been deployed
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
3 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
Adds:
priorattribute to parameters.with_log_priorfunction which takes an arbitraryObjectiveand adds the log prior of the parameters of the model to it. This allows for MAP estimation of parameters, rather than solely MLE.