Streaming variational inference: Trainer for minibatch ADVI

[YichengYang-Ethan]

Streaming variational inference: Trainer for minibatch ADVI

Stacks on #698 (the streaming DataLoader). This ports the Trainer from
pymc-devs/pymc#8333 into extras, so both halves of the streaming-VI work —
out-of-core data plus a callback-free fit loop — are here now.

What it does

Trainer(method="advi", dataloader=loader, data_name="batch").fit(n) owns the
fitting loop and streams each minibatch into the model's pm.Data placeholder
with set_data, so the user writes no callbacks. The DataLoader owns batching
(len(dataloader) is the dataset size N), the model owns the math.

from pymc_extras.variational.streaming import DataLoader, Trainer, parquet_source

loader = DataLoader(
    parquet_source("shuffled/"), batch_size=4096, sample_shape=(4,), total_size="auto"
)
with pm.Model() as model:
    b = pm.Normal("b", 0.0, 3.0, shape=4)
    batch = pm.Data("batch", np.zeros((4096, 4)))
    logit = b[0] + b[1] * batch[:, 0] + b[2] * batch[:, 1] + b[3] * batch[:, 2]
    pm.Bernoulli("y", logit_p=logit, observed=batch[:, 3], total_size=len(loader))
    approx = Trainer(method="advi", dataloader=loader, data_name="batch").fit(20_000)

fit(n) feeds exactly n minibatches (the first seeds the placeholder before
step 0; the advance after the final step is skipped). User callbacks compose with
the internal advance instead of colliding on the keyword, and an Inference
instance is forwarded to pm.fit unchanged.

Relationship to #635 (the new ADVI API)

#635 already reworks ADVI around its own Trainer
(pymc_extras/inference/advi/training.py), so this is not meant to land a
second, competing Trainer. It's the interim, out-of-core path that runs on
today's pm.fit — useful now, and a concrete home for the streaming DataLoader
while #635 is in review. The end state is to feed the DataLoader into #635's
Trainer through a small adapter, at which point this interim Trainer retires.
Opening it as a draft so the two can be reconciled rather than duplicated; happy
to fold it into #635's direction once that settles.

Where the scaling lives (interim)

The N / batch_size rescaling stays in the model, via total_size=len(loader),
reusing the existing create_minibatch_rv machinery, so it runs unchanged on
today's pm.fit. Folding the scaling into the inference step — so it is derived
from len(dataloader) and drops out of the model body — is the cleaner end state
and lines up with #635; the Trainer Notes document this.

Tests

tests/variational/test_streaming_trainer.py: end-to-end equivalence to in-RAM
pm.Minibatch ADVI, exact batch accounting (fit(n) consumes n batches),
placeholder seeding/streaming, refine resuming the stream, the pass-boundary
total_size check, user-callback composition, and the input guards.

Refs

Ports pymc-devs/pymc#8333 · stacks on #698 · original data layer
pymc-devs/pymc#8325 · relates to #635.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.86%. Comparing base (86fac3c) to head (b736877).
⚠️ Report is 50 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #710       +/-   ##
===========================================
+ Coverage   51.60%   91.86%   +40.26%     
===========================================
  Files          73      101       +28     
  Lines        8003     9159     +1156     
===========================================
+ Hits         4130     8414     +4284     
+ Misses       3873      745     -3128     

Files with missing lines	Coverage Δ
pymc_extras/variational/__init__.py	100.00% <100.00%> (ø)
pymc_extras/variational/trainer.py	100.00% <100.00%> (ø)

... and 31 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]

Is this ready for review?

[YichengYang-Ethan]

Marking this ready, with a short reading guide.

trainer.py top to bottom:

• _cycle — the only epoch logic, and it lives here rather than in the loader on purpose: the loader yields exactly one epoch per iteration and knows nothing about repetition. _cycle restarts it forever and turns an empty source into a raise instead of an infinite spin. Happy to inline it into fit if you'd rather not have the helper.
• _warn_if_scaling_mismatches — the one guard I kept, because both failure modes are silent: a model that never declares total_size, or declares a different one than the loader streams, fits successfully and just gives a wrong posterior. Per Monday's discussion it warns and never blocks — the user may know something we don't.
• Trainer.fit — cycle the loader, seed the placeholder with batch 0, advance one batch per step through a callback appended after the user's. Step i trains on batch i; user callbacks run while the batch that produced the latest loss is still installed, and refine() continues the stream instead of repeating a batch.

Tests pin the batch-to-step alignment against the loader itself, the duck-typed loader contract (what the deleted isinstance guard used to cover), and both mismatch warnings. Validated end to end on 30.7M rows of tick data against a closed-form posterior streamed through the same loader.