Streaming variational inference: Trainer for minibatch ADVI

[YichengYang-Ethan]

Streaming variational inference: Trainer for minibatch ADVI

Stacks on #698 (the streaming DataLoader). This adds the Trainer that
drives variational inference over a DataLoader with no user-facing
callbacks — the second half of the streaming-VI work, ported here from
pymc-devs/pymc#8333 now that the data layer lives in extras.

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, and the advance after the final step is skipped. User callbacks (e.g. a
convergence tracker) compose with the internal advance instead of colliding on
the keyword. An Inference instance is forwarded to pm.fit unchanged.

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. This runs on today's
pm.fit unchanged. Folding the scaling into the inference step — so the
Trainer derives it from len(dataloader) and it drops out of the model body —
is the cleaner end state and fits naturally with the VI rework
(fit_advi / Inference.step(batch)); it will follow once that inference-step
API lands. The Notes section of Trainer documents 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.