New ADVI API by jessegrabowski · Pull Request #635 · pymc-devs/pymc-extras

jessegrabowski · 2026-02-02T03:22:09Z

This PR moves the work from pymc-devs/pymc#7799 over here to extras. The key idea is to copy the numpyro guide model API, but with our own PyMC flair.

I also added:

stick the landing estimator
forward sampling helper
LLM sketch for a training API patterned after pytorch-lightning

I updated the notebook to use the proposed training API. It obviously needs a lot of work (being that it's llm trash) but hopefully it can get some ideas flowing.

review-notebook-app · 2026-02-02T03:22:15Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

codecov-commenter · 2026-02-02T04:53:10Z

Codecov Report

❌ Patch coverage is 93.42857% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.99%. Comparing base (86fac3c) to head (380120a).
⚠️ Report is 53 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_extras/inference/advi/training.py	93.37%	11 Missing ⚠️
pymc_extras/inference/advi/optimizers.py	92.36%	10 Missing ⚠️
pymc_extras/inference/advi/compile.py	96.55%	1 Missing ⚠️
pymc_extras/inference/advi/fit.py	94.44%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #635       +/-   ##
===========================================
+ Coverage   51.60%   92.99%   +41.38%     
===========================================
  Files          73      107       +34     
  Lines        8003     9749     +1746     
===========================================
+ Hits         4130     9066     +4936     
+ Misses       3873      683     -3190

Files with missing lines	Coverage Δ
pymc_extras/inference/advi/__init__.py	`100.00% <100.00%> (ø)`
pymc_extras/inference/advi/autoguide.py	`98.11% <100.00%> (+98.11%)`	⬆️
pymc_extras/inference/advi/compile.py	`97.77% <96.55%> (ø)`
pymc_extras/inference/advi/fit.py	`94.44% <94.44%> (ø)`
pymc_extras/inference/advi/optimizers.py	`92.36% <92.36%> (ø)`
pymc_extras/inference/advi/training.py	`93.41% <93.37%> (+93.41%)`	⬆️

... and 44 files with indirect coverage changes

🚀 New features to boost your workflow:

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

zaxtax

I think this is pretty neat! After some rebasing and small API changes we should definitely try to get this merged!

zaxtax · 2026-02-13T16:54:40Z

+        The probabilistic model.
+    guide : AutoGuideModel
+        The variational guide.
+    stick_the_landing : bool, optional


I feel like we need a better name than stick_the_landing . Also the function is supposed to return the logp and logq terms but the STL estimator is about returning only the path derivative component of the gradient.

zaxtax · 2026-02-13T17:02:33Z

+            Number of MC draws per step for gradient estimation, by default 10.
+        model : Model
+            The PyMC model to fit. If None, the model is inferred from context.
+        state : SVIState, optional


I'd prefer we copy what PyTorch Lightning did here and have the SVI State live in the Trainer object and be something we can pass upon initalisation.

zaxtax · 2026-02-13T17:03:01Z

+        return state
+
+    def sample_posterior(
+        self, draws: int, state: SVIState, model: Model | None = None


drop SVIState and move into the Trainer object

ricardoV94 · 2026-06-13T16:48:55Z

+    model: Model | None = None,
+    *,
+    n_steps: int = 10_000,
+    draws_per_step: int = 8,


Suggested change

draws_per_step: int = 8,

draws_per_step: int = 1,

fonnesbeck · 2026-06-18T21:12:02Z

Is the plan to deprecate the OPVI fit function when this merged? (or matriculates to the main project?)

zaxtax · 2026-06-18T22:15:37Z

            Deterministic(
                rv.name,
-                loc + pt.softplus(scale) * z,
+                loc + pt.exp(scale) * z,


We had it with softplus since that's what numpyro does? Why change?

2x faster, logdet becomes -scale, instead of log(softplus)

It's not a sure choice, just something I'm exploring

ricardoV94 · 2026-06-18T22:23:27Z

Is the plan to deprecate the OPVI fit function when this merged? (or matriculates to the main project?)

There's no big plan on my end, still very much WIP. Maybe if it proves itself. The VI part of this is def smaller in scope (but also complexity)

jessegrabowski · 2026-06-18T23:38:44Z

Is the plan to deprecate the OPVI fit function when this merged? (or matriculates to the main project?)

From my part, this refactor was an effort to explore a different design. Could maybe lead to a full replacement of pm.fit, but the scope so far is just ADVI. I'm not sure the particle methods would be in scope for what we have in this PR.

But I think the guide model setup is really interesting, and works really well with the pymc Model API + pytensor graph mutations. So its appealing in and of itself.

ricardoV94 · 2026-07-05T17:04:01Z

@zaxtax let's leave the minibatch out of this. I already had quite some discussion with @jessegrabowski I don't think PyMC models have the right API for us to eagerly go and guess what's batched data and what's not, and assume it's over the 0-th axis and so forth. Check the over-enginereered #700 I started for what I think is the minimum work a library built on top of PyMC would have to do to "auto minibatch". I would not put any minibatch stuff on this PR. For now if users want to minibatch they'd use the lover-or-hate it pm.Minibatch path

zaxtax · 2026-07-05T17:07:28Z

But the proposed API here offers a high degree of control, so I'm not sure what's ambiguous here

…

On Sun, 5 Jul 2026, 19:04 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> @zaxtax <https://github.com/zaxtax> let's leave the minibatch out of this. I already had quite some discussion with @jessegrabowski <https://github.com/jessegrabowski> I don't think PyMC models have the right API for us to eagerly go and guess what's batched data and what's not, and assume it's over the 0-th axis and so forth. Check the over-enginereered #700 <#700> I started for what I think is the minimum work a library built on top of PyMC would have to do to "auto minibatch". I would not put any minibatch stuff on this PR. For now if users want to minibatch they'd use the lover-or-hate it pm.Minibatch path — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUKW4QZRK4RZFOTJPST5DKDBRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYGY4DGNRYHA42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4886836889>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUKXGY54EK4L2RNIXZL5DKDBRAVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUIZRBNMLPDGTEJL4ID5DKDBRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYGY4DGNRYHA42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUMAMPL6EIZPB5SPODL5DKDBRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYGY4DGNRYHA42M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

ricardoV94 · 2026-07-05T18:51:21Z

But the proposed API here offers a high degree of control, so I'm not sure what's ambiguous here

That's exactly the sort of discussion I didn't want to have here... Anyway I see no place to control which X variables you stream and how. Your proposal is a bit more low-level and less debatable than an "auto minibatch", but I still would prefer to leave it to its own PR. The scaling thing is also a bit naive (in terms of calling model.logp() per var which can be costly). I didn't check how it handles univariate vs multivariates. Also why merging observed transform into the fitting routine?

zaxtax · 2026-07-05T19:00:16Z

I'm happy to have that be part of a separate PR. I'll try to on that future PR really go into detail about the approach.

…

On Sun, 5 Jul 2026, 20:51 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> But the proposed API here offers a high degree of control, so I'm not sure what's ambiguous here I don't want to have the discussion, but I see no place to control which X variables you stream and how. Your proposal is a bit more low-level and less debatable than an "auto minibatch", but I still would prefer to leave it to its own PR. The scaling thing is also a bit naive (in terms of calling model.logp() per var which can be costly). I didn't check how it handles univariate vs multivariates. Also why merging observed transform into the fitting routine? — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUKYOOIVYCJWGXJOVGT5DKPT5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYG4ZDCMJVGAZKM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4887211502>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULT7JQVBMWKJQGF6IL5DKPT5AVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUKBM2UO5ZI4UOITUND5DKPT5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYG4ZDCMJVGAZKM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCULNG4HJD2BTWKMNJLL5DKPT5A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYG4ZDCMJVGAZKM4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

zaxtax · 2026-07-17T01:12:06Z

Did my Trainer class commit end up in another PR?

…

On Fri, 17 Jul 2026, 03:08 Jesse Grabowski, ***@***.***> wrote: *jessegrabowski* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> @ricardoV94 <https://github.com/ricardoV94> can I close this PR as superseded by the other ones? — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUKXOJ44B4ASKYYMVG35FF363A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJZG44TSNRYG4Y2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4997996871>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUN3APYV32IKZDLYW635FF363AVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUK5OUIN53MNDMTN7V35FF363A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJZG44TSNRYG4Y2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUKXJ2FOGG5YIMWIPD35FF363A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJZG44TSNRYG4Y2M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

ricardoV94 · 2026-07-17T13:28:45Z

@ricardoV94 can I close this PR as superseded by the other ones?

no, but this needs rebasing of things that went elsewhere

ricardoV94 · 2026-07-17T13:30:30Z

Did my Trainer class commit end up in another PR?

I reverted the minibatch commit, the idea was to discuss in another PR yes. I didn't open yet

zaxtax · 2026-07-17T13:39:46Z

There was another Trainer commit in addition to that one around minibatch

…

On Fri, 17 Jul 2026, 15:30 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> Did my Trainer class commit end up in another PR? I reverted the minibatch commit, the idea was to discuss in another PR yes. I didn't open yet — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUPZP4V7ZSQ6D5QSHZD5FITAZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM3TQOBRGE22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5003788115>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULPYYWNQTFAELZS5SL5FITAZAVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUJH2DNVZEJMMRW44ET5FITAZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM3TQOBRGE22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUOXTF5IFKK3HNZLSOL5FITAZA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM3TQOBRGE22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

ricardoV94 · 2026-07-17T13:40:42Z

I tried to be careful to only revert, not delete stuff that wasn't mine. But may have failed

zaxtax · 2026-07-17T13:40:50Z

I mean this one c186229

zaxtax · 2026-07-17T13:44:35Z

But the Trainer class I don't think is a major design change in my opinion at least

…

On Fri, 17 Jul 2026, 15:41 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> I tried to be careful to only revert, not delete stuff that wasn't mine. But may have failed — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUJU72BP4MTQWJYX7JD5FIUHBA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM4DOMZWGI22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5003873625>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUJXXNNI5IVGKTCYWMD5FIUHBAVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUNJMFMXA4VC7Y6O6AT5FIUHBA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM4DOMZWGI22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUK3PJNP24WRIMDED2T5FIUHBA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGM4DOMZWGI22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

ricardoV94 · 2026-07-17T13:51:57Z

I'll take a look, but I suspect we may have 3 chefs who want 3 slight different things here.

Personally, I'd distinguish fit machinery state and the fit result. For the fit machinery state we should try to reuse the pymc step stampler state api/machinery (not the step samplers, those are different things). https://github.com/pymc-devs/pymc/blob/df1a369d5fbdc0cc49082b0c0ef7cbd511cf76e9/pymc/step_methods/state.py#L54

We can reject if we find a good reason why it doesn't fit the mold. It doesn't need to happen in this PR either, but I'd like to plan on checking it.

Or data structures / classes, I don't want to incuur overhead of storing each step result in a fancy dataclass or promise a state that was stale (i.e. read the shared variables content on the first iteration and pretend those are up to date in the pure pytensor path). I don't know if I pushed the latest state of the api notebook, but most my decisions were based on bringing inference speed down, and they checked out there (in the last example brought from @juanitorduz blogpost). When the inner step function gets cheap I don't want to be losing to jax/numpyro.

My one criteria is I don't want @juanitorduz to have an excuse to use numpyro for SVI. Or at least not a good one.

On the control over the fit loop, and what we expose vs not I don't have much opinion, but both you @jessegrabowski and @zaxtax need to check if you're on the same page. I didn't look into that at all but may have accidentally reverted or erased something (sorry if erased, it wasn't supposed to if it happened)

zaxtax · 2026-07-17T14:08:14Z

I'll take a look, but I suspect we may have 3 chefs who want 3 slight different things here.

Personally, I'd distinguish fit machinery state and the fit result. For the fit machinery state we should try to reuse the pymc step stampler state api/machinery (not the step samplers, those are different things). https://github.com/pymc-devs/pymc/blob/df1a369d5fbdc0cc49082b0c0ef7cbd511cf76e9/pymc/step_methods/state.py#L54

We can reject if we find a good reason why it doesn't fit the mold. It doesn't need to happen in this PR either, but I'd like to plan on checking it.

Or data structures / classes, I don't want to incuur overhead of storing each step result in a fancy dataclass or promise a state that was stale (i.e. read the shared variables content on the first iteration and pretend those are up to date in the pure pytensor path). I don't know if I pushed the latest state of the api notebook, but most my decisions were based on bringing inference speed down, and they checked out there (in the last example brought from @juanitorduz blogpost). When the inner step function gets cheap I don't want to be losing to jax/numpyro.

My one criteria is I don't want @juanitorduz to have an excuse to use numpyro for SVI. Or at least not a good one.

On the control over the fit loop, and what we expose vs not I don't have much opinion, but both you @jessegrabowski and @zaxtax need to check if you're on the same page. I didn't look into that at all but may have accidentally reverted or erased something (sorry if erased, it wasn't supposed to if it happened)

I don't have a strong opinion here but I worry sampler state is a poor fit for the different state representations ADVI methods can have. As for the fit loop, all I ask is @jessegrabowski have a peek about what I made in the existing commit (possibly the reverted one) and we confirm if we are on the same page. If we are not, I'm happy for us to just pick something and go with it

juanitorduz · 2026-07-17T14:08:17Z

I will be happy to test this against a real-life example https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb :D.

In the meantime,I wanna try to run this (as in, the current api, not this one) on GPU + Modal (I made the numpyro version fit this model in less than 10 min with SVI and 50K time series), https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb is the same example, and we tried to port the batching machinery pymc-labs/pymc-forecast#48

zaxtax · 2026-07-17T14:23:35Z

And yes I want @juanitorduz to love whatever we make

juanitorduz · 2026-07-17T14:25:02Z

And yes I want @juanitorduz to love whatever we make

I AM loving it ... that is why we are pushing with @twiecki a testing field : https://github.com/pymc-labs/pymc-forecast :D

ricardoV94 · 2026-07-17T15:00:04Z

I will be happy to test this against a real-life example https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb :D.

In the meantime,I wanna try to run this (as in, the current api, not this one) on GPU + Modal (I made the numpyro version fit this model in less than 10 min with SVI and 50K time series), https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb is the same example, and we tried to port the batching machinery pymc-labs/pymc-forecast#48

what batching machinery? minibatch or vectorization -adjacent?

juanitorduz · 2026-07-17T17:54:07Z

I will be happy to test this against a real-life example https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb :D.
In the meantime,I wanna try to run this (as in, the current api, not this one) on GPU + Modal (I made the numpyro version fit this model in less than 10 min with SVI and 50K time series), https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb is the same example, and we tried to port the batching machinery pymc-labs/pymc-forecast#48

what batching machinery? minibatch or vectorization -adjacent?

For posterior predictive sampling. In the NumPyro function Predictive() this was killing the memory, and the "trick" that worked was to move this computation to the CPU and do it by batches.

zaxtax · 2026-07-17T18:00:27Z

Oh this is a good insight. Can PyMC PPC be run in batches over the data?

…

On Fri, 17 Jul 2026, 19:54 Juan Orduz, ***@***.***> wrote: *juanitorduz* left a comment (pymc-devs/pymc-extras#635) <#635 (comment)> I will be happy to test this against a real-life example https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb :D. In the meantime,I wanna try to run this (as in, the current api, not this one) on GPU + Modal (I made the numpyro version fit this model in less than 10 min with SVI and 50K time series), https://github.com/pymc-labs/pymc-forecast/blob/main/docs/examples/retail_stockouts.ipynb is the same example, and we tried to port the batching machinery pymc-labs/pymc-forecast#48 <pymc-labs/pymc-forecast#48> what batching machinery? minibatch or vectorization -adjacent? For posterior predictive sampling. In the NumPyro function Predictive() this was killing the memory, and the "trick" that worked was to move this computation to the CPU and do it by batches. — Reply to this email directly, view it on GitHub <#635?email_source=notifications&email_token=AAACCUNN5CYIPW2DE7LYMUL5FJR5LA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGU4TQNRRGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5005986157>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUL47EI6QPE2YTPZZLL5FJR5LAVCNFSNUABFKJSXA33TNF2G64TZHM2DGNBWG43TAOJTHNEXG43VMU5TGOBYGM3TGMZSGAYKC5QC> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUJWLPUZW7AG5W33ZDT5FJR5LA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGU4TQNRRGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUJJU764N7NLBDLQFYT5FJR5LA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMBQGU4TQNRRGU32M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you were mentioned.Message ID: ***@***.***>

zaxtax · 2026-08-13T10:21:58Z

+        self.convergence_window = convergence_window
+        self.relative_tolerance = relative_tolerance
+
+        self.model = model


Should we make Trainer object modelcontext aware?

zaxtax · 2026-08-13T10:22:25Z

+        return self._guide
+
+    @property
+    def clip_norm(self) -> float | None:


Do we need this?

zaxtax · 2026-08-13T10:24:07Z

Do we need learning rate schedulers in this PR?

zaxtax · 2026-08-13T10:24:43Z

+    assert not np.array_equal(draws_a, draws_c)
+
+
+def test_fit_advi_early_stopping(conjugate_model):


Do we need early stopping support in this PR?

zaxtax · 2026-08-20T15:10:35Z

@ricardoV94 what's left to do on this one other than maybe remove learning rate and clip norm from the Trainer object?

fit_advi entry point with an optax-like optimizer API and numpyro-style defaults, a compiled SVI step that bakes in draws and runs optimizer updates (Adam bias-correction kept in floatX), the training loop, and the objective.

Applies to the mean-field, full-rank and low-rank guides, so the scale parameterization is the same across all three.

Replace the SVIModule/ADVIModule/SVITrainer split with one Trainer that owns the training loop, following the design of pymc-devs/pymc#8333 and PyTorch Lightning: all configuration (guide, optimizer, learning rate, convergence-based early stopping, model, backend) lives at construction, there are no user-facing hooks or callbacks, and fit(n) just runs. The duplicate fit/fit_jitted loops are merged into a single fit that internally picks the compiled fast path (clipped Adam baked into the step function; the default) or the Python-side update loop (when an optax-like optimizer is passed). The fast path now also supports resuming parameters from a passed SVIState. The trainer keeps the last state, so sample_posterior() works without arguments. fit_advi becomes a thin wrapper over Trainer, and draws_per_step is renamed to n_particles and moved to the constructor.

fit(n) now always continues from the current parameters and Adam moments, replacing a state=None argument that meant "start fresh" on the python path and "silently resume" on the compiled one. reset(), load_state() and the state property cover those cases explicitly, and SVIState carries the Adam moments so resuming in another trainer matches continuing in place. State is read out of the shared variables once at the end of a fit rather than rebuilt per step, so the loop no longer allocates a state object it cannot fill in honestly. Learning rates are resolved up front for the whole call, lifting the schedule (and its np.interp) out of the hot loop. Configuration splits the same way: what gets compiled into the step function is read-only, what is per-run policy is read afresh by each fit call, with a learning_rate override for a single call. Drops the python-side optimizer loop, which cost 61us/step against 3.6us for the compiled path, was used by no test, notebook or default, and was where both divergences came from. optimizers.py loses the unreachable optax-like transformations and becomes schedules.py. fit_advi is unchanged bit-for-bit.

Cut it from 76 cells to 26, dropping the SVML/llvmlite probes, the timeit scratch cells, the dead SGD and RMSProp optimizer classes, and the duplicate benchmark runs. What remains is fit_advi and Trainer on three models: a linear regression against NUTS, the Trainer API itself (fit continuing, state snapshots, load_state, reset, per-call learning rate), radon with an LKJ covariance, and the 116k-parameter forecasting benchmark. The radon comparison is a two-panel scatter over all 85 counties rather than an 85-row forest plot: same information, and 3.4MB less of it. The benchmark runs 10_000 steps instead of 3_000. At 3_000 the one-cycle rate has fully annealed while sigma is still descending, which read as fit_advi losing to pm.fit on quality when it was only losing on step count. All three now land on sigma 0.313 and seasonal correlation 0.876. Also drops the claim that numba loses to XLA here. On pytensor catch_up_with_jax at 4be21562a the ADVI step is 6.7ms on numba against 7.7ms on jax; what is left is numba's cold compile, which the notes now say instead.

fit(n) now runs n steps. The window check interacted badly with the default schedule: the one-cycle is sized to the requested n, so stopping early exits wherever that schedule had reached, which for an early stop is near the peak. On radon it fired at 2800 of 10000 steps at lr 0.00706 against a 0.008 peak, i.e. at close to maximum step size. Spending the same 2800 steps with the schedule sized to 2800 lands fully annealed and gets 2.2x lower mean error against a nutpie reference (0.0122 vs 0.0273), so the heuristic was worse than simply asking for fewer steps. That leaves n as a pacing parameter and not just a cap, which the docstrings now say. Deciding when to stop by hand is better served by what the trainer already does: fit in chunks, look at state.loss_history, and call fit again to continue from the same parameters and Adam moments. KeyboardInterrupt also keeps the state, so an interrupted run is resumable. Schedules are now functions of the trainer's global step rather than of the offset within a call, and the default one-cycle is sized to start_step + n. A follow-up fit anneals the rest of one cycle instead of ramping a second time, which was the last thing about resuming that did not carry over. Resuming from a snapshot still matches continuing in place, since load_state restores the step count the schedule reads. Removes convergence_window and relative_tolerance from Trainer and fit_advi, and reruns the notebook.

…imizers, add sgd Port the optimizer refactor from advi-minibatch into advi-refactor. - Add optimizers.py with GradientTransformation (init/update/pytensor), adam, clipped_adam, sgd, rmsprop, chain, clip_by_global_norm, and linear_onecycle_schedule; remove schedules.py. - compile_svi_step_fn now takes an optimizer and returns (step_fn, shared_params, shared_optimizer_state); step_fn takes no inputs. - Trainer takes an optimizer instead of learning_rate/clip_norm and defaults to clipped_adam(); fit() no longer accepts a per-call learning_rate. - fit_advi takes an optimizer instead of learning_rate/clip_norm. - Preserve the optimizer-state snapshot/resume feature: SVIState.optimizer_state is read from the shared variables, so state/load_state/reset keep working (empty for stateless optimizers like sgd).

jessegrabowski added enhancements New feature or request help wanted Extra attention is needed inference labels Feb 2, 2026

jessegrabowski requested review from AlexAndorra, Dekermanjian, lucianopaz, ricardoV94 and zaxtax February 2, 2026 03:22

zaxtax mentioned this pull request Feb 2, 2026

Move VI autoguide work from pymc repo #636

Merged

zaxtax reviewed Feb 13, 2026

View reviewed changes

ricardoV94 force-pushed the advi-refactor branch 3 times, most recently from 6525424 to a661c84 Compare June 12, 2026 16:29

ricardoV94 reviewed Jun 13, 2026

View reviewed changes

ricardoV94 force-pushed the advi-refactor branch 3 times, most recently from 68e83fe to eefecfb Compare June 16, 2026 14:23

zaxtax reviewed Jun 18, 2026

View reviewed changes

ricardoV94 mentioned this pull request Jun 28, 2026

Build autoguide in unconstrained space and add full and low rank multivariate guides #701

Merged

ricardoV94 force-pushed the advi-refactor branch from 091ff86 to 14a94a4 Compare July 6, 2026 10:38

YichengYang-Ethan mentioned this pull request Jul 12, 2026

Streaming variational inference: Trainer for minibatch ADVI #710

Closed

ricardoV94 force-pushed the advi-refactor branch from 14a94a4 to 525d9e1 Compare July 19, 2026 21:29

ricardoV94 mentioned this pull request Aug 13, 2026

Reworking Minibatching for ADVI #713

Merged

zaxtax approved these changes Aug 13, 2026

View reviewed changes

ricardoV94 force-pushed the advi-refactor branch from 9bc08d5 to 96224fb Compare August 14, 2026 15:19

ricardoV94 and others added 7 commits August 26, 2026 20:53

Add ADVI fit API with compiled SVI step and optimizers

14ebc23

fit_advi entry point with an optax-like optimizer API and numpyro-style defaults, a compiled SVI step that bakes in draws and runs optimizer updates (Adam bias-correction kept in floatX), the training loop, and the objective.

Use same backend for deterministics

0921355

Use exp unconstraining in the ADVI guides

b7141c2

Applies to the mean-field, full-rank and low-rank guides, so the scale parameterization is the same across all three.

jessegrabowski force-pushed the advi-refactor branch from 96224fb to d14cb71 Compare August 27, 2026 02:01

zaxtax merged commit d48872a into pymc-devs:main Aug 27, 2026
29 checks passed

		assert not np.array_equal(draws_a, draws_c)


		def test_fit_advi_early_stopping(conjugate_model):

Uh oh!

Conversation

jessegrabowski commented Feb 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

review-notebook-app Bot commented Feb 2, 2026

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

zaxtax 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fonnesbeck commented Jun 18, 2026

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 commented Jun 18, 2026

Uh oh!

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

Uh oh!

Uh oh!

ricardoV94 commented Jul 5, 2026

Uh oh!

zaxtax commented Jul 5, 2026 via email

Uh oh!

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

Uh oh!

Uh oh!

zaxtax commented Jul 5, 2026 via email

Uh oh!

zaxtax commented Jul 17, 2026 via email

Uh oh!

ricardoV94 commented Jul 17, 2026

Uh oh!

ricardoV94 commented Jul 17, 2026

Uh oh!

zaxtax commented Jul 17, 2026 via email

Uh oh!

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

Uh oh!

Uh oh!

zaxtax commented Jul 17, 2026

Uh oh!

zaxtax commented Jul 17, 2026 via email

Uh oh!

ricardoV94 commented Jul 17, 2026

Uh oh!

zaxtax commented Jul 17, 2026

Uh oh!

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

Uh oh!

Uh oh!

zaxtax commented Jul 17, 2026

Uh oh!

juanitorduz commented Jul 17, 2026

Uh oh!

ricardoV94 commented Jul 17, 2026

Uh oh!

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

Uh oh!

Uh oh!

zaxtax commented Jul 17, 2026 via email

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

jessegrabowski commented Feb 2, 2026 •

edited

Loading

codecov-commenter commented Feb 2, 2026 •

edited

Loading

jessegrabowski commented Jun 18, 2026 •

edited

Loading

ricardoV94 commented Jul 5, 2026 •

edited

Loading

ricardoV94 commented Jul 17, 2026 •

edited

Loading

juanitorduz commented Jul 17, 2026 •

edited

Loading

juanitorduz commented Jul 17, 2026 •

edited

Loading