Repository navigation
Conversation
bwengals
added this pull request to stack #70
October 10, 2026 06:36
…mpute Inducing points are the unique rows of X, each conditioning on its k nearest preceding points, so the data and KL terms of the ELBO minibatch independently. The variational covariance is block diagonal (block_size=1 is mean-field). Neighbors are found host side with cKDTree or a correlation metric, held in shared variables, and recomputed between optimizer steps; prediction finds test-point neighbors through predict_spec and refuses to run if the hyperparameters the metric depends on have changed since the last recompute. compile_training_step gains extra_inputs, and compile_diagnostics evaluates a diagnostics namedtuple against trained parameters. The cubic-floor counter now unwraps BlockwiseWithCoreShape, which the Numba backend uses for batched ops, and joint_graph_analysis.py reads the scalar field of the objective namedtuples.
…inputs, skip frozen vars in grad_norm pg.fit cannot pass the index inputs vnngp_elbo needs, so VNNGP no longer advertises a default objective. The correlation-metric kernel evaluator accepts float64 inputs under float32 floatX, and compile_diagnostics takes frozen_vars so grad_norm covers only trained parameters.
bwengals
marked this pull request as ready for review
October 15, 2026 23:12
jessegrabowski
approved these changes
Oct 16, 2026
| ## Where things live | ||
|
|
||
| - **Models** — `ptgp/gp/`. `Unapproximated`, `VFE`, `SVGP`. Each exposes | ||
| - **Models** — `ptgp/gp/`. `Unapproximated`, `VFE`, `SVGP`, `VNNGP`. Each exposes |
Member
There was a problem hiding this comment.
VNNGP has no predict_joint
|
|
||
| MLLTerms = namedtuple("MLLTerms", ["mll", "fit", "logdet"]) | ||
| ELBOTerms = namedtuple("ELBOTerms", ["elbo", "var_exp", "kl"]) | ||
| VNNGPELBOTerms = namedtuple("VNNGPELBOTerms", ["elbo", "var_exp", "kl"]) |
Member
There was a problem hiding this comment.
same fields as ELBOTerms, return that
| b : tensor, shape (B, K) | ||
| F : tensor, shape (B,) | ||
| """ | ||
| from pytensor.graph.replace import vectorize_graph |
|
|
||
| def _rows_match(X, row_idx, value): | ||
| """Return ``value``, asserting at run time that ``X`` and ``row_idx`` have equal length.""" | ||
| from pytensor.raise_op import Assert |
| Global norm of the training-loss gradient on this batch. Filled in by | ||
| :func:`ptgp.optim.compile_diagnostics`; NaN when evaluated directly. | ||
| """ | ||
| from ptgp.likelihoods import op_of |
| ls = pt.dscalar("ls") | ||
| gp = Unapproximated(kernel=ExpQuad(input_dim=1, ls=ls), mean=Zero(), sigma=sigma) | ||
| loss = marginal_log_likelihood(gp, X, y) | ||
| loss = marginal_log_likelihood(gp, X, y).mll |
Member
There was a problem hiding this comment.
separate PR against main, the script raises AttributeError there today
| def _op_name(op): | ||
| return type(op.core_op).__name__ if isinstance(op, Blockwise) else type(op).__name__ | ||
| # Blockwise, and the BlockwiseWithCoreShape wrapper the Numba backend uses, expose core_op | ||
| return type(getattr(op, "core_op", op)).__name__ |
Member
There was a problem hiding this comment.
same PR as scripts/joint_graph_analysis.py:122
| op = node.op | ||
| core = op.core_op if isinstance(op, Blockwise) else op | ||
| # Blockwise, and the BlockwiseWithCoreShape wrapper the Numba backend uses, expose core_op | ||
| core = getattr(node.op, "core_op", node.op) |
Member
There was a problem hiding this comment.
same PR as scripts/joint_graph_analysis.py:122
| return predict_fn | ||
|
|
||
|
|
||
| def compile_diagnostics( |
Member
There was a problem hiding this comment.
move to a follow-up after #62, with vnngp_diagnostics and the pitfall pages
| ) | ||
|
|
||
|
|
||
| def vnngp_diagnostics(vnngp, X, y, row_idx, kl_idx): |
Member
There was a problem hiding this comment.
follow-up, same as ptgp/optim/training.py:1332
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
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.
VNNGP suits large, low-dimensional data such as spatial or spatiotemporal fields, where the function has local structure on lengthscales much shorter than the domain. Placing an inducing point at every input lets it resolve detail that a few thousand global inducing points cannot, and in Wu et al. (Table 1) it beats SVGP clearly on UKHousing, 3DRoad and Precipitation. SVGP remains the better choice for smooth functions of many inputs, where a low-rank approximation is efficient and SVGP matched or beat VNNGP on several higher-dimensional datasets (Elevators, Bike).
Merge order: #66 (
predict_spec), then #68 (likelihood graph, replaces #50), then this PR. It is stacked on #68 and needs retargeting to main once both merge.Adds
VNNGP(Wu, Pleiss & Cunningham, ICML 2022, arXiv:2202.01694). Inducing points are the unique rows ofX, and each conditions on itsknearest preceding points in a random or maximin ordering, so both ELBO terms minibatch at O(k³) per point.ptgp/neighbors.py: orderings, neighbor search (cKDTree onX / scale, or the correlation metric of Kang & Katzfuss 2023), and block partitioning.block_size=1is mean-field.recompute_neighborsrefreshes neighbors under trained hyperparameters, such as ARD lengthscales, between optimizer steps. Prediction finds test-point neighbors throughpredict_specand raises if the hyperparameters changed since the last recompute.vnngp_elbo(vnngp, X, y, row_idx, kl_idx)takes separate data and KL minibatches.compile_training_stepgainsextra_inputsfor the indices.vnngp_diagnosticsand a genericcompile_diagnostics, which also reports the training-loss gradient norm for any diagnostics namedtuple with agrad_normfield. Pitfall pages underdocs/agents/ptgp/pitfalls/follow the format of Rework the agent skill into a general GP diagnosis skill #62.(batch, k, k). The cubic-floor counter now unwrapsBlockwiseWithCoreShape, which the Numba backend uses, andscripts/joint_graph_analysis.pyis fixed for the namedtuple objectives.pg.fitdoes not supportVNNGPyet; that follows oncepg.fitmoves to pytensor-ml's L-BFGS.📚 Documentation preview 📚: https://ptgp--69.org.readthedocs.build/en/69/