VFF implementation - #2
Closed
thomaspinder wants to merge 61 commits into
Closed
Conversation
Project setup stuff
…nary and non kernels, just use ipynb not md tracking
Bernoulli and Poisson are cross-checked against GPJax; StudentT and NegativeBinomial against a self-contained numpy+scipy reference ELBO (whitened predict by hand, scipy.integrate.quad for the variational expectation, closed-form whitened KL), since GPJax does not ship those likelihoods. Both patterns pin the full ELBO wiring at atol=1e-5.
…redict_marginal The existing predict returns per-point marginals, which can't be used to draw smooth function samples — neighboring points' uncertainty is uncorrelated. predict_joint gives the full (N, N) posterior covariance, and predict_f_samples Cholesky-transforms caller-supplied iid-normal noise into draws from that joint. This is the foundation for posterior-predictive use with non-Gaussian likelihoods: sample f, push through the likelihood. Renamed predict to predict_marginal across SVGP, VFE, Unapproximated to make the pair symmetric and the semantics explicit. base_conditional gains a full_cov flag.
…iag onto kernels Deletes ~300 lines of unused JAX-based fitting code and its re-exports, along with two empty test files. Replaces the kernel_diag method on each GP model with a diag(X) method on each kernel, written as a direct expression (ones for stationary/Gibbs, x for RandomWalk, delegated for WarpedInput, composed for Sum/Product) rather than diagonal of the full matrix. PyTensor's rewrite system can't derive this automatically because it requires semantic knowledge of the kernel.
…mplement _eval Each kernel previously repeated the Y-is-None branch and the specify_assumptions(symmetric, positive_definite) wrap. That scaffolding now lives once in Kernel.__call__; subclasses provide a pure _eval(X, Y) that's always called with both arguments. Shrinks every kernel by several lines and leaves one place to change the annotation policy.
These classes only support a Gaussian likelihood, so expose the noise parameter directly and build the Gaussian object internally.
Unapproximated and VFE stored _X_train/_y_train attributes that were never assigned; VFE.predict_marginal even read self._X_train without initializing it. Make X_train/y_train required args on both predict_marginal methods. tests/conftest.py declared X_1d/X_2d/rng_key fixtures and referenced a make_params helper that never existed; the fixtures were shadowed by local definitions in every test that used those names.
Both hardcoded input_dim=1, preventing their use in sum/product kernels and with multi-column inputs. Accept any input_dim; enforce a single active column via active_dims (defaulting to [0] when input_dim == 1).
Follows PyMC's pattern (pm.sample, pm.sample_prior_predictive): a single compile_kwargs dict forwarded to pytensor.function, covering mode (NUMBA/JAX/C), allow_input_downcast, on_unused_input, etc. Applied to compile_training_step, compile_scipy_objective, compile_predict, and greedy_variance_init. compile_training_step merges any caller-provided updates with the optimizer's updates (mirroring pymc.pytensorf.compile).
Covers POSITIVE propagation (Sqr/Pow/Mul/DimShuffle/Alloc/ExtractDiag), the new POSITIVE_DEFINITE inference paths (AllocDiag of symbolic positive vector, Mul of positive scalar × PSD, set_subtensor diagonal pattern, PSD-through-transpose, X.T·M·X and X.T·Solve(M,X) quadratic forms), the specify_assumptions(positive=True) signature, and end-to-end SLogDet lowering (firing, non-firing, Cholesky-sharing, numeric correctness). Test bodies use only pytensor APIs so they can be lifted into pytensor when the rules are upstreamed; the only ptgp imports are the side-effect registration and the POSITIVE key.
Two scale-unlock features for SVGP on big-N data, planned as independent PRs: 1D-Matérn variational Fourier features via a new InducingVariables subclass with a Woodbury solve_Kuu bridge, and an exact closed-form variational expectation for Poisson with exp link. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Four findings addressed: - PR 2 (Poisson): the original plan proposed adding an analytic VE that already exists in ptgp/likelihoods/poisson.py. Rescoped to hardening (end-to-end SVGP test, gradient correctness, graph-level regression assertion). Dropped the invlink-callable→link-string API change, which would have been a silent breaking change for users passing custom invlinks. - PR 1 (VFF) domain validation: specified runtime checks for training and prediction, a default refusal to extrapolate, and an opt-in allow_extrapolation escape hatch with explicit tests. - PR 1 (VFF) Kuu contract: split K_uu (dense) from structured_Kuu (optional typed DiagPlusLowRank). solve_Kuu is the only dispatch point, so a buggy consumer calling K_uu on a VFF object gets a slow but correct result, not a shape-incompatible structured tuple. - Algorithmic-path tests: both PRs gain CI-asserted graph-introspection tests that verify the fast path actually runs in the compiled graph, rather than relying on wall-clock notebooks. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
|
I screwed something up and didnt mean to close this... will fix. Sorry about that! |
Collaborator
|
ok looks like I cant reopen because I force pushed to main and now the git history isnt shared. Working on opening a new PR from a new banch, and then cherry picking your commits to that so you dont have to worry about doing a rebase |
Collaborator
|
ok so I tried that, and it looks like it/I am making a huge mess. My fault for pushing straight to main... Will start using branches properly now. Would you mind opening a new PR or doing a rebase or whatevers easier? |
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.
An implementation of Variational Fourier Features for Gaussian Processes. The implementation is based upon a fork of the authors' original code here. The main new functionality lies within the kernel's inducing fourier points, whilst the remaining code is predominantly wiring modules together and ensuring back compatibility e.g.,
base_conditionalbeing connected to a whitened and unwhitened form.On documentation, I have run a notebook locally that tests for equality in this implementation and the original authors' code. I did not commit though as it muddies the dependencies. I have extended the demo notebook to include the VFF approximation, and also to quantify the distance from each approximation's posterior to the exact posterior using the Wasserstein-2 distance (closed form when you have MVNormals).