Build autoguide in unconstrained space and add full and low rank multivariate guides by ricardoV94 · Pull Request #701 · pymc-devs/pymc-extras

ricardoV94 · 2026-06-26T19:16:50Z

Took the unconstrain fix from #635 and added some multivariate guides for fun

codecov-commenter · 2026-06-27T11:47:33Z

Codecov Report

❌ Patch coverage is 96.75325% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.40%. Comparing base (86fac3c) to head (76e4dff).
⚠️ Report is 18 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_extras/inference/advi/autoguide.py	96.85%	4 Missing ⚠️
pymc_extras/inference/advi/training.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #701       +/-   ##
===========================================
+ Coverage   51.60%   91.40%   +39.80%     
===========================================
  Files          73       92       +19     
  Lines        8003     8543      +540     
===========================================
+ Hits         4130     7809     +3679     
+ Misses       3873      734     -3139

Files with missing lines	Coverage Δ
pymc_extras/inference/advi/__init__.py	`100.00% <100.00%> (ø)`
pymc_extras/inference/advi/compile.py	`100.00% <100.00%> (ø)`
pymc_extras/inference/advi/objective.py	`84.61% <100.00%> (+84.61%)`	⬆️
pymc_extras/inference/advi/pytensorf.py	`100.00% <100.00%> (+100.00%)`	⬆️
pymc_extras/inference/advi/training.py	`0.00% <0.00%> (ø)`
pymc_extras/inference/advi/autoguide.py	`97.48% <96.85%> (+97.48%)`	⬆️

... and 28 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2026-06-28T14:35:47Z

+    if "trust_input" not in compile_kwargs:
+        compile_kwargs["trust_input"] = True


use set_default

zaxtax

I love this

ricardoV94 · 2026-06-28T16:27:24Z

+            preserves_shape = (
+                transform is None
+                or transform.ndim_supp == 0
+                or equal_computations([value_shape], [rv_shape])


this is a stop gag until we have transform.ndim_supp in: pymc-devs/pymc#8331

is anything blocking that? I approved it.

pytensor being broken, will likely only ship with next release

ricardoV94 · 2026-06-28T16:27:57Z

-        """Returns a graph representing the logp of the guide model, evaluated under draws from its random variables."""
+    @property
+    def latent(self) -> Variable:
+        """The whole unconstrained draw, before it is split into per-variable values.


probably makes sense for the regular Normal to also do it?

Should this be a private or public method?

I don't think it matters? I have no strong preference, let me know if you do

jessegrabowski

looks nice, gave suggestions

jessegrabowski · 2026-06-30T13:33:55Z

+    lower-triangular Cholesky factor. The guide's logq is obtained by logprob inference over the
+    un-split draw (see :meth:`AutoFullRankGuideModel.stochastic_logq`); ``L`` is annotated
+    lower-triangular via :func:`pytensor.assumptions.assume` so the internal ``solve``/``slogdet``
+    lower to the triangular fast path.


Internal computation details not important to end-user, this is a public facing function so I'd rather the docstring be usage focused

jessegrabowski · 2026-06-30T13:34:16Z

+    value_shapes, value_dims = get_value_shapes_and_dims(model)
+
+    # DictToArrayBijection flattens the unconstrained initial point into the guide mean init;
+    # point_map_info gives the per-RV layout (hence the order) of that flat vector.


Suggested change

# point_map_info gives the per-RV layout (hence the order) of that flat vector.

# point_map_info gives the per-RV layout (thus the order) of that flat vector.

jessegrabowski · 2026-06-30T13:40:37Z

+        # cholesky outputs are auto-recognized as lower-triangular by pytensor's linalg rewrites
+        chol_cap = cholesky(capacitance, lower=True)
+        rhs = W.T @ (delta * d_inv)
+        sol = solve_triangular(chol_cap, rhs, lower=True)


this should maybe be cho_solve, need to benchmark it but it's always been close when i tried it (or check that we rewrite this case. I know it would end up as cho_solve if this was a basic solve, not sure about solve_triangular)

jessegrabowski · 2026-06-30T13:41:10Z

+
+        delta = u - loc
+        d_inv = 1.0 / d**2
+        capacitance = pt.eye(rank) + (W * d_inv[:, None]).T @ W  # K x K, symmetric PD


"Capacitance is the ability of an object to store electric charge when a voltage is applied across it. It is measured in farads (F) and depends on factors like the size of the conductive plates, the distance between them, and the type of dielectric material used."

Why this name?

Due to a manifold collapse in high dimensional next token prediction that we'll never be able to fully phantom and yet must respect if not enjoy.

This is what the bot had to say in its defense

It traces back to the "capacitance matrix method" for solving elliptic PDEs on irregular domains (Buzbee–Dorr–George–Golub, 1971), by analogy to capacitance in circuit network theory — the low-rank correction plays the role of a boundary/capacitance term. It carried into general Woodbury usage. Concretely, PyTorch's LowRankMultivariateNormal uses this exact name (_capacitance_tril), so it's the conventional term for precisely this distribution.

jessegrabowski · 2026-06-30T13:43:24Z

+    if rank is None:
+        rank = round(n_dim**0.5)
+    rank = max(1, min(rank, n_dim))


Not sure I like this. I'd rather force users to think about it than give a random default with no real motivation

the motivation is that numpyro uses this default. hard against?

review-notebook-app · 2026-07-06T07:15:46Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

Construct AutoDiagonalNormal guides over the model's unconstrained value variables, matching the standard-normal parameter dtype to the value dtype to avoid a per-element cast in the numba RV core loop.

ricardoV94 force-pushed the mvn_autogide branch from 39af9c1 to a774967 Compare June 27, 2026 11:41

ricardoV94 force-pushed the mvn_autogide branch 2 times, most recently from b65846b to 559e951 Compare June 27, 2026 13:13

ricardoV94 changed the title ~~Multivariate autoguide~~ Build autoguide in unconstrained space and add full and low rank guides Jun 28, 2026

ricardoV94 force-pushed the mvn_autogide branch from 559e951 to 30b0c3f Compare June 28, 2026 14:34

ricardoV94 requested review from jessegrabowski and zaxtax June 28, 2026 14:34

ricardoV94 commented Jun 28, 2026

View reviewed changes

ricardoV94 changed the title ~~Build autoguide in unconstrained space and add full and low rank guides~~ Build autoguide in unconstrained space and add full and low rank multivariate guides Jun 28, 2026

ricardoV94 added bug Something isn't working enhancements New feature or request inference labels Jun 28, 2026

zaxtax approved these changes Jun 28, 2026

View reviewed changes

Comment thread pymc_extras/inference/advi/autoguide.py Outdated

ricardoV94 commented Jun 28, 2026

View reviewed changes

ricardoV94 force-pushed the mvn_autogide branch from 30b0c3f to f103e29 Compare June 30, 2026 09:29

jessegrabowski approved these changes Jun 30, 2026

View reviewed changes

ricardoV94 force-pushed the mvn_autogide branch from f103e29 to 8cb114c Compare July 6, 2026 07:15

ricardoV94 force-pushed the mvn_autogide branch 3 times, most recently from 30b0c3f to f4b17e2 Compare July 6, 2026 09:51

Build ADVI autoguides in the unconstrained value space

69193fa

Construct AutoDiagonalNormal guides over the model's unconstrained value variables, matching the standard-normal parameter dtype to the value dtype to avoid a per-element cast in the numba RV core loop.

ricardoV94 force-pushed the mvn_autogide branch from f4b17e2 to 888675d Compare July 6, 2026 09:52

Add low rank and fullrank mvnormal autoguides

76e4dff

ricardoV94 force-pushed the mvn_autogide branch from 888675d to 76e4dff Compare July 6, 2026 10:20

jessegrabowski merged commit f4f46ab into pymc-devs:main Jul 10, 2026
29 checks passed

ricardoV94 mentioned this pull request Jul 14, 2026

New ADVI fails for variables whose value transforms changes the variables' shape #646

Closed

		if "trust_input" not in compile_kwargs:
		compile_kwargs["trust_input"] = True

	# point_map_info gives the per-RV layout (hence the order) of that flat vector.
	# point_map_info gives the per-RV layout (thus the order) of that flat vector.

Uh oh!

Conversation

ricardoV94 commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov-commenter commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

review-notebook-app Bot commented Jul 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ricardoV94 commented Jun 26, 2026 •

edited

Loading

codecov-commenter commented Jun 27, 2026 •

edited

Loading

ricardoV94 Jul 6, 2026 •

edited

Loading