From 048868ee5a06bcfbdda69b4b0a58779ee31ac268 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 18 Jul 2026 22:06:12 -0500 Subject: [PATCH 1/9] Add streaming Trainer for minibatch ADVI Trainer(method=..., dataloader=...).fit(n) owns the loop: it seeds the model's pm.Data placeholder before step 0 and streams a batch into it after every step, so the user writes no callbacks. Every step advances, including the last. Skipping the final advance made fit(n) pull exactly n batches, but it left the batch fit had just trained in the placeholder, so Inference.refine -- which steps before replaying callbacks -- retrained it, and the closure counter that implemented the skip stayed live inside refine after an early stop and stranded it on a stale batch. The loader already reads one batch ahead for its pass-size check, so uniform advancing costs nothing a re-readable source did not already pay. User callbacks run before the advance, so one inspecting the placeholder sees the batch that produced the latest loss rather than its successor, and a StopIteration from one ends the fit without pulling again. An Inference instance bound to a different model than the one being trained is now refused instead of silently optimizing a model that never receives a batch, and a model whose observed variables declare no total_size, or one that disagrees with the loader's N, warns rather than returning a quietly misweighted posterior. --- pymc_extras/variational/__init__.py | 2 + pymc_extras/variational/trainer.py | 198 ++++++++++++++++ tests/variational/test_trainer.py | 344 ++++++++++++++++++++++++++++ 3 files changed, 544 insertions(+) create mode 100644 pymc_extras/variational/trainer.py create mode 100644 tests/variational/test_trainer.py diff --git a/pymc_extras/variational/__init__.py b/pymc_extras/variational/__init__.py index 782872e7c..620d33d57 100644 --- a/pymc_extras/variational/__init__.py +++ b/pymc_extras/variational/__init__.py @@ -16,9 +16,11 @@ parquet_source, shuffle_buffer, ) +from pymc_extras.variational.trainer import Trainer __all__ = [ "DataLoader", + "Trainer", "parquet_source", "shuffle_buffer", ] diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py new file mode 100644 index 000000000..df3b55366 --- /dev/null +++ b/pymc_extras/variational/trainer.py @@ -0,0 +1,198 @@ +# Copyright 2026 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Drive variational inference over a :class:`DataLoader` with no user callbacks.""" + +from __future__ import annotations + +import warnings + +from collections.abc import Iterator + +import numpy as np + +from pymc.model import modelcontext +from pymc.variational.inference import Inference +from pymc.variational.inference import fit as _fit +from pymc.variational.minibatch_rv import MinibatchRandomVariable +from pytensor.graph.basic import Constant + +from pymc_extras.variational.dataloader import DataLoader, _is_positive_int + +__all__ = ["Trainer"] + + +def _warn_if_scaling_mismatches(model, total_size: int) -> None: + """Warn if the model cannot rescale the minibatch likelihood to ``total_size``. + + ``total_size`` reaches the graph as a constant on a ``MinibatchRandomVariable``. + Both an absent and a disagreeing value are silent otherwise: the posterior comes + out wrong rather than the fit failing. + """ + declared = { + int(t.data) + for rv in model.observed_RVs + if isinstance(rv.owner.op, MinibatchRandomVariable) + for t in rv.owner.inputs[1:] + if isinstance(t, Constant) and t.data is not None + } + if not declared: + warnings.warn( + "no observed variable declares total_size, so the minibatch " + "log-likelihood is never rescaled and is underweighted against the prior " + "by about N / batch_size. Pass total_size=len(dataloader) to the observed " + "distribution.", + UserWarning, + stacklevel=3, + ) + elif declared != {total_size}: + warnings.warn( + f"the model declares total_size={sorted(declared)} but the loader streams " + f"N={total_size}; the minibatch rescaling does not match the data.", + UserWarning, + stacklevel=3, + ) + + +class Trainer: + """Drive variational inference over a :class:`DataLoader` without user callbacks. + + Follows the design in PyMC's variational-inference rework and PyTorch + Lightning: the ``Trainer`` owns the training loop, the + :class:`DataLoader` owns batching (and ``len(dataloader)`` is the dataset size + ``N``), and the model owns the math. The model exposes a ``pm.Data`` placeholder; + the ``Trainer`` streams minibatches into it with ``model.set_data`` once per + step; no user callbacks are needed. + + Parameters + ---------- + method : str or Inference, default "advi" + Variational method, forwarded to :func:`pymc.fit`: a name (``"advi"``, + ``"fullrank_advi"``, ...) or an :class:`~pymc.variational.inference.Inference` + instance. ``pm.fit`` applies ``model`` and ``random_seed`` only to a name; + an instance is already bound to a model, so configure it at construction + (e.g. ``ADVI(random_seed=...)``). + dataloader : DataLoader + The minibatch source. ``len(dataloader)`` is ``N``; the model should pass + it to the observed distribution's ``total_size``. + model : pymc.Model, optional + Defaults to the model on the context stack. + data_name : str, default "batch" + Name of the ``pm.Data`` placeholder minibatches are streamed into. Must + match the name used for ``pm.Data(name, ...)`` in the model. + **fit_kwargs + Default keyword arguments forwarded to :func:`pymc.fit` (e.g. + ``obj_optimizer``); per-call kwargs to :meth:`fit` override them. + + Notes + ----- + The per-step ``set_data`` currently lives in the ``Trainer``. Once the VI + rework's ``Inference.step(batch)`` lands it moves there, at which point the + ``total_size`` rescaling can be derived from ``len(dataloader)`` and dropped + from the model body entirely. + + Examples + -------- + .. code-block:: python + + 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))) # placeholder + 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) + """ + + def __init__( + self, + *, + method: str | Inference = "advi", + dataloader: DataLoader, + model=None, + data_name: str = "batch", + **fit_kwargs, + ): + self.method = method + self.dataloader = dataloader + self.model = model + self.data_name = data_name + self._fit_kwargs = fit_kwargs + + def fit(self, n: int = 10_000, **kwargs): + """Fit for ``n`` steps, streaming minibatches into the model's placeholder. + + Step ``i`` trains on batch ``i``: the first batch seeds the placeholder + before step 0 and every step loads the next one, so ``n`` steps train ``n`` + batches and leave batch ``n`` loaded for whatever runs next -- + :meth:`~pymc.variational.inference.Inference.refine` then continues the + stream instead of repeating a batch. User ``callbacks`` run while the batch + that produced the latest loss is still in place, and a ``StopIteration`` + from one ends the fit before another batch is loaded. Keyword arguments are + forwarded to :func:`pymc.fit` on top of the constructor's ``fit_kwargs`` + (per-call wins); ``progressbar`` defaults to ``False`` unless either sets it. + + Returns + ------- + :class:`Approximation` + The fitted approximation, as returned by :func:`pymc.fit`. + """ + if not _is_positive_int(n): + raise ValueError(f"n must be a positive integer (the number of fit steps), got {n!r}") + loader = self.dataloader + if not isinstance(loader, DataLoader): + raise TypeError( + f"Trainer needs a DataLoader for `dataloader`, got {type(loader).__name__}." + ) + model = modelcontext(self.model) + if self.data_name not in model: + raise KeyError( + f"data_name {self.data_name!r} is not a variable in the model; it " + f"must name the pm.Data placeholder the minibatches are streamed into." + ) + if isinstance(self.method, Inference) and self.method.approx.model is not model: + raise ValueError( + "`method` is an Inference instance bound to a different model than the " + "one being trained, so the minibatches would stream into a model the " + "fit never reads. Build it under this model, or pass a method name." + ) + if loader.total_size is not None: + _warn_if_scaling_mismatches(model, loader.total_size) + + def _stream() -> Iterator[np.ndarray]: + while True: + empty = True + for batch in loader: + empty = False + yield batch + if empty: + raise RuntimeError("dataloader yielded no batches") + + batches = _stream() + model.set_data(self.data_name, next(batches)) + + def _advance(*_): + model.set_data(self.data_name, next(batches)) + + merged = {**self._fit_kwargs, **kwargs} + merged.setdefault("progressbar", False) + user_callbacks = merged.pop("callbacks", None) or [] + return _fit( + n, + method=self.method, + model=model, + callbacks=[*user_callbacks, _advance], + **merged, + ) diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py new file mode 100644 index 000000000..eec1f8e44 --- /dev/null +++ b/tests/variational/test_trainer.py @@ -0,0 +1,344 @@ +# Copyright 2026 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Trainer: drive variational inference over a DataLoader with no user callbacks.""" + +import numpy as np +import pymc as pm +import pytest + +from pymc_extras.variational.dataloader import DataLoader +from pymc_extras.variational.trainer import Trainer +from tests.variational.dataloader_helpers import chunked_factory + + +def marked(n, rows=4): + """``n`` blocks, block ``i`` filled with ``i``, so an installed batch names itself.""" + return [np.full((rows, 1), float(i)) for i in range(n)] + + +def record_installed(model, log): + """Log the marker of every batch ``set_data`` installs, in order.""" + original = model.set_data + + def spy(name, values, *args, **kwargs): + log.append(float(np.asarray(values)[0, 0])) + return original(name, values, *args, **kwargs) + + model.set_data = spy + + +def test_trainer_end_to_end_matches_in_ram_minibatch(): + """End-to-end: Trainer-driven streaming ADVI reproduces in-RAM pm.Minibatch ADVI. + + Exercises the whole API: a pm.Data placeholder, total_size=len(loader), and a + Trainer that streams minibatches into the placeholder with set_data while the + user writes no callbacks. Runs long enough to cycle the loader across epochs. + """ + seed = 0 + rng = np.random.default_rng(seed) + N, bs = 60_000, 2048 + X = rng.normal(size=(N, 2)) + b_true = np.array([0.3, -1.1, 0.7]) + y = (rng.random(N) < 1 / (1 + np.exp(-(b_true[0] + X @ b_true[1:])))).astype("float64") + data = np.column_stack([X, y]) + + with pm.Model(): + b = pm.Normal("b", 0, 3, shape=3) + xb, zb, yb = pm.Minibatch(X[:, 0].copy(), X[:, 1].copy(), y, batch_size=bs) + pm.Bernoulli("o", logit_p=b[0] + b[1] * xb + b[2] * zb, observed=yb, total_size=N) + ap = pm.fit( + 6000, + method="advi", + obj_optimizer=pm.adam(learning_rate=0.02), + progressbar=False, + random_seed=seed, + ) + in_ram = ap.sample(400).posterior["b"].values.reshape(-1, 3).mean(0) + + loader = DataLoader( + chunked_factory(data, 20_000), + batch_size=bs, + shuffle=True, + buffer_size=40_000, + seed=seed, + sample_shape=(3,), + total_size=N, + ) + with pm.Model() as model: + b = pm.Normal("b", 0, 3, shape=3) + batch = pm.Data("batch", np.zeros((bs, 3))) + pm.Bernoulli( + "o", + logit_p=b[0] + b[1] * batch[:, 0] + b[2] * batch[:, 1], + observed=batch[:, 2], + total_size=len(loader), + ) + ap = Trainer( + method="advi", + dataloader=loader, + data_name="batch", + obj_optimizer=pm.adam(learning_rate=0.02), + ).fit(6000, random_seed=seed) + stream = ap.sample(400).posterior["b"].values.reshape(-1, 3).mean(0) + + np.testing.assert_allclose(in_ram, stream, atol=0.1) + + +def test_trainer_streams_into_placeholder(): + """The Trainer seeds the pm.Data placeholder before step 0 (pm.fit runs + callbacks after each step) and overwrites it each step; after fitting it holds + a real batch, not the zero seed.""" + data = np.ones((4, 1)) + loader = DataLoader(lambda: iter([data] * 100), batch_size=4, sample_shape=(1,), total_size=4) + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + Trainer(method="advi", dataloader=loader, data_name="batch").fit( + 5, progressbar=False, random_seed=0 + ) + np.testing.assert_array_equal(model["batch"].get_value(), data) + + +def test_trainer_raises_when_loader_cannot_restart(): + """A source that streams one epoch and then comes back empty cannot be cycled; + the Trainer surfaces a clear error instead of training on stale data.""" + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + if calls["n"] == 1: + yield np.zeros((4, 1)) + + loader = DataLoader(factory, batch_size=4, sample_shape=(1,), total_size=4) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + with pytest.raises(RuntimeError, match="yielded no batches"): + Trainer(method="advi", dataloader=loader, data_name="batch").fit( + 5, progressbar=False, random_seed=0 + ) + + +def test_trainer_rejects_non_dataloader(): + """The isinstance guard fires before any model lookup.""" + with pytest.raises(TypeError, match="DataLoader"): + Trainer(method="advi", dataloader=object()).fit(10) + + +def test_trainer_appends_user_callbacks_and_streams_distinct_batches(): + """User callbacks (e.g. convergence trackers) compose with the internal + advance callback instead of colliding on the keyword, and the placeholder + holds a different batch on successive steps. Also exercises the default + data_name ("batch").""" + blocks = [np.full((4, 1), float(i)) for i in range(60)] + loader = DataLoader(lambda: iter(blocks), batch_size=4, sample_shape=(1,), total_size=240) + seen = [] + with pm.Model() as model: + x = pm.Normal("x", 0.0, 1.0) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", x, 1.0, observed=batch[:, 0], total_size=len(loader)) + Trainer(method="advi", dataloader=loader).fit( + 5, callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))] + ) + assert len(seen) == 5 + assert len(set(seen)) > 1 + + +def test_trainer_accepts_inference_instance(): + """An Inference instance is forwarded to pm.fit unchanged; it is bound to + the model it was built under, so the Trainer only streams the batches.""" + data = np.ones((4, 1)) + loader = DataLoader(lambda: iter([data] * 50), batch_size=4, sample_shape=(1,), total_size=4) + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + approx = Trainer(method=pm.ADVI(random_seed=0), dataloader=loader).fit(5) + assert len(approx.hist) == 5 + np.testing.assert_array_equal(model["batch"].get_value(), data) + + +def test_constructor_fit_kwargs_take_random_seed(): + """random_seed works as a constructor default, as the docstring promises, + and a per-call value overrides the constructor's.""" + data = np.ones((4, 1)) + + def fit_with(ctor_kwargs, fit_kwargs): + loader = DataLoader( + lambda: iter([data] * 50), batch_size=4, sample_shape=(1,), total_size=4 + ) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + return Trainer(method="advi", dataloader=loader, data_name="batch", **ctor_kwargs).fit( + 5, **fit_kwargs + ) + + a = fit_with({"random_seed": 7}, {}) + b = fit_with({"random_seed": 0}, {"random_seed": 7}) + np.testing.assert_array_equal(a.hist, b.hist) + + +def test_fit_trains_one_batch_per_step(): + """Step i trains batch i, and the step that ends the fit loads batch n for what follows.""" + loader = DataLoader( + lambda: iter(marked(10, rows=2)), batch_size=2, sample_shape=(1,), total_size=20 + ) + installed = [] + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((2, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + record_installed(model, installed) + Trainer(method="advi", dataloader=loader).fit(3, random_seed=0) + assert installed == [0.0, 1.0, 2.0, 3.0] + + +def test_refine_after_fit_continues_without_repeating_a_batch(): + """Inference.refine replays pm.fit's saved callbacks and steps before they run. + + fit therefore has to leave the *next* batch loaded, or refine's first gradient + step would retrain the batch fit just finished with. + """ + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + installed = [] + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + record_installed(model, installed) + inference = pm.ADVI(random_seed=0) + Trainer(method=inference, dataloader=loader).fit(3) + assert installed == [0.0, 1.0, 2.0, 3.0] + installed.clear() + inference.refine(4, progressbar=False) + assert installed == [4.0, 5.0, 6.0, 7.0] + + +def test_refine_after_an_early_stop_keeps_streaming(): + """A fit cut short by a callback must not leave the advance permanently disarmed.""" + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + installed = [] + armed = [True] + + def stop_once(*_): + if armed[0]: + armed[0] = False + raise StopIteration("stop") + + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + record_installed(model, installed) + inference = pm.ADVI(random_seed=0) + Trainer(method=inference, dataloader=loader).fit(10, callbacks=[stop_once], score=False) + assert installed == [0.0] + installed.clear() + inference.refine(4, progressbar=False) + assert installed == [1.0, 2.0, 3.0, 4.0] + + +def test_user_callbacks_see_the_batch_that_produced_the_loss(): + """A callback reading the placeholder must see batch i on step i, not batch i+1.""" + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + seen = [] + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + Trainer(method="advi", dataloader=loader).fit( + 5, + random_seed=0, + callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))], + ) + assert seen == [0.0, 1.0, 2.0, 3.0, 4.0] + + +def test_inference_instance_bound_to_another_model_is_rejected(): + """Streaming into one model while an Inference optimizes another is silent otherwise.""" + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + with pm.Model() as other: + pm.Normal("mu", 0, 1) + pm.Data("batch", np.zeros((4, 1))) + elsewhere = pm.ADVI() + assert other is not None + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + with pytest.raises(ValueError, match="bound to a different model"): + Trainer(method=elsewhere, dataloader=loader).fit(3) + + +@pytest.mark.parametrize( + "declared, match", + [(None, "no observed variable declares total_size"), (40, "does not match the data")], + ids=["absent", "wrong-value"], +) +def test_unusable_likelihood_scaling_warns(declared, match): + """Absent or disagreeing total_size biases the posterior without failing the fit.""" + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=declared) + with pytest.warns(UserWarning, match=match): + Trainer(method="advi", dataloader=loader).fit(3, random_seed=0) + + +def test_total_size_check_fires_when_fit_ends_at_pass_boundary(): + """fit(n) with n exactly the batches in one pass still runs the total_size + sanity check: the stream is kept one batch ahead, so stopping at the + boundary does not abandon the check right before it would fire.""" + data = np.zeros((40, 1)) + loader = DataLoader(chunked_factory(data, 10), batch_size=10, sample_shape=(1,), total_size=400) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((10, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + with pytest.warns(UserWarning, match="disagrees with"): + Trainer(method="advi", dataloader=loader).fit(4, random_seed=0) + + +def test_fit_rejects_nonpositive_n(): + """fit consumes the seed batch before pm.fit could reject n itself, so a + non-positive n is refused up front, before touching the stream.""" + loader = DataLoader( + lambda: iter([np.zeros((2, 1))]), batch_size=2, sample_shape=(1,), total_size=2 + ) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((2, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + with pytest.raises(ValueError, match="positive integer"): + Trainer(method="advi", dataloader=loader).fit(0) + + +def test_unknown_data_name_raises_before_consuming(): + """A data_name that is not in the model raises a guided KeyError before any + batch is pulled from the loader.""" + loader = DataLoader( + lambda: iter([np.zeros((4, 1))] * 3), batch_size=4, sample_shape=(1,), total_size=4 + ) + installed = [] + with pm.Model() as model: + pm.Normal("mu", 0, 1) + record_installed(model, installed) + with pytest.raises(KeyError, match=r"pm\.Data placeholder"): + Trainer(method="advi", dataloader=loader, data_name="nope").fit(2) + assert installed == [] From bc97d8b04c568c0c7cd99ccef541d59083b64d16 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 1 Aug 2026 18:29:16 -0500 Subject: [PATCH 2/9] Pin the Trainer's batch-to-step alignment against the loader itself Decodes which batch each gradient step trained from the loss fingerprint and compares it to the sequence a user gets by iterating the loader -- across a single step, a fit ending exactly at an epoch seam, and fits wrapping the loader once and several times. The previous assertions checked that n distinct values were seen, which a stream off by one still satisfies. Also pins that the scaling warning fires on an absent or mismatched total_size and stays quiet on a correct one, over two N/batch pairs, and that the bound-model check accepts an instance built under the model as readily as it rejects one built elsewhere. --- tests/variational/test_trainer.py | 187 ++++++++++++++++++++++++++---- 1 file changed, 164 insertions(+), 23 deletions(-) diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index eec1f8e44..4f6937d9b 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -13,6 +13,10 @@ # limitations under the License. """Trainer: drive variational inference over a DataLoader with no user callbacks.""" +import warnings + +from itertools import chain, islice, repeat + import numpy as np import pymc as pm import pytest @@ -27,6 +31,16 @@ def marked(n, rows=4): return [np.full((rows, 1), float(i)) for i in range(n)] +def loud(n, rows=4, spacing=1000.0): + """``n`` blocks, block ``i`` filled with ``spacing * (i + 1)``. + + Markers this far from zero and from each other survive into the loss, which is + quadratic in the batch value, so the loss fingerprints the batch the gradient + was actually taken on rather than the one ``set_data`` was asked for. + """ + return [np.full((rows, 1), spacing * (i + 1)) for i in range(n)] + + def record_installed(model, log): """Log the marker of every batch ``set_data`` installs, in order.""" original = model.set_data @@ -38,6 +52,11 @@ def spy(name, values, *args, **kwargs): model.set_data = spy +def streamed(loader, k): + """The first ``k`` markers a user iterating ``loader`` themselves gets, epoch after epoch.""" + return [float(batch[0, 0]) for batch in islice(chain.from_iterable(repeat(loader)), k)] + + def test_trainer_end_to_end_matches_in_ram_minibatch(): """End-to-end: Trainer-driven streaming ADVI reproduces in-RAM pm.Minibatch ADVI. @@ -72,7 +91,6 @@ def test_trainer_end_to_end_matches_in_ram_minibatch(): shuffle=True, buffer_size=40_000, seed=seed, - sample_shape=(3,), total_size=N, ) with pm.Model() as model: @@ -100,7 +118,7 @@ def test_trainer_streams_into_placeholder(): callbacks after each step) and overwrites it each step; after fitting it holds a real batch, not the zero seed.""" data = np.ones((4, 1)) - loader = DataLoader(lambda: iter([data] * 100), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter([data] * 100), batch_size=4, total_size=4) with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -121,7 +139,7 @@ def factory(): if calls["n"] == 1: yield np.zeros((4, 1)) - loader = DataLoader(factory, batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(factory, batch_size=4, total_size=4) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -144,7 +162,7 @@ def test_trainer_appends_user_callbacks_and_streams_distinct_batches(): holds a different batch on successive steps. Also exercises the default data_name ("batch").""" blocks = [np.full((4, 1), float(i)) for i in range(60)] - loader = DataLoader(lambda: iter(blocks), batch_size=4, sample_shape=(1,), total_size=240) + loader = DataLoader(lambda: iter(blocks), batch_size=4, total_size=240) seen = [] with pm.Model() as model: x = pm.Normal("x", 0.0, 1.0) @@ -161,7 +179,7 @@ def test_trainer_accepts_inference_instance(): """An Inference instance is forwarded to pm.fit unchanged; it is bound to the model it was built under, so the Trainer only streams the batches.""" data = np.ones((4, 1)) - loader = DataLoader(lambda: iter([data] * 50), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter([data] * 50), batch_size=4, total_size=4) with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -177,9 +195,7 @@ def test_constructor_fit_kwargs_take_random_seed(): data = np.ones((4, 1)) def fit_with(ctor_kwargs, fit_kwargs): - loader = DataLoader( - lambda: iter([data] * 50), batch_size=4, sample_shape=(1,), total_size=4 - ) + loader = DataLoader(lambda: iter([data] * 50), batch_size=4, total_size=4) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -195,9 +211,7 @@ def fit_with(ctor_kwargs, fit_kwargs): def test_fit_trains_one_batch_per_step(): """Step i trains batch i, and the step that ends the fit loads batch n for what follows.""" - loader = DataLoader( - lambda: iter(marked(10, rows=2)), batch_size=2, sample_shape=(1,), total_size=20 - ) + loader = DataLoader(lambda: iter(marked(10, rows=2)), batch_size=2, total_size=20) installed = [] with pm.Model() as model: mu = pm.Normal("mu", 0, 1) @@ -208,13 +222,55 @@ def test_fit_trains_one_batch_per_step(): assert installed == [0.0, 1.0, 2.0, 3.0] +@pytest.mark.parametrize( + "n, blocks, rows", + [(1, 6, 4), (4, 4, 4), (5, 3, 2), (7, 2, 1)], + ids=["single-step", "ends-at-seam", "wraps-once", "wraps-often"], +) +def test_steps_consume_the_loaders_own_batch_sequence(n, blocks, rows): + """The batches trained are the loader's own pass, in order, across epoch seams. + + Three views have to agree with what iterating the loader by hand yields: the + batches ``set_data`` installed, the placeholder a callback reads while the loss + it was handed is still current, and the losses themselves -- with the optimizer + frozen the loss is quadratic in the batch marker, so ``hist / hist[0]`` decodes + the batch each gradient was taken on. A seam that repeats, drops or reorders a + batch still passes a "``n`` distinct batches" check. + """ + loader = DataLoader( + lambda: iter(loud(blocks, rows=rows)), + batch_size=rows, + total_size=blocks * rows, + ) + expected = streamed(loader, n + 1) + installed, seen = [], [] + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((rows, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + record_installed(model, installed) + approx = Trainer( + method="advi", + dataloader=loader, + obj_optimizer=pm.sgd(learning_rate=1e-12), + ).fit( + n, + random_seed=0, + callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))], + ) + assert installed == expected + assert seen == expected[:n] + trained = np.asarray(expected[:n]) / expected[0] + np.testing.assert_allclose(approx.hist / approx.hist[0], trained**2, rtol=0.02) + + def test_refine_after_fit_continues_without_repeating_a_batch(): """Inference.refine replays pm.fit's saved callbacks and steps before they run. fit therefore has to leave the *next* batch loaded, or refine's first gradient step would retrain the batch fit just finished with. """ - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) installed = [] with pm.Model() as model: mu = pm.Normal("mu", 0, 1) @@ -231,7 +287,7 @@ def test_refine_after_fit_continues_without_repeating_a_batch(): def test_refine_after_an_early_stop_keeps_streaming(): """A fit cut short by a callback must not leave the advance permanently disarmed.""" - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) installed = [] armed = [True] @@ -253,9 +309,37 @@ def stop_once(*_): assert installed == [1.0, 2.0, 3.0, 4.0] +def test_a_second_fit_reseeds_and_forgets_the_first_calls_kwargs(): + """Calling fit twice on one Trainer: each call re-reads the constructor's defaults + and installs a batch before its own step 0. + + A per-call keyword that survived into the next call would silently change a fit + nobody passed it to, and a call that trusted the batch its predecessor left + loaded would train it twice. + """ + loader = DataLoader(lambda: iter(marked(6)), batch_size=4, total_size=24) + installed, seen = [], [] + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + record_installed(model, installed) + trainer = Trainer(method="advi", dataloader=loader, random_seed=0) + unscored = trainer.fit(3, score=False) + left_loaded = installed[-1] + installed.clear() + scored = trainer.fit( + 3, callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))] + ) + assert len(unscored.hist) == 0 + assert len(scored.hist) == 3 + assert seen == installed[:-1] + assert installed[0] != left_loaded + + def test_user_callbacks_see_the_batch_that_produced_the_loss(): """A callback reading the placeholder must see batch i on step i, not batch i+1.""" - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) seen = [] with pm.Model() as model: mu = pm.Normal("mu", 0, 1) @@ -271,7 +355,7 @@ def test_user_callbacks_see_the_batch_that_produced_the_loss(): def test_inference_instance_bound_to_another_model_is_rejected(): """Streaming into one model while an Inference optimizes another is silent otherwise.""" - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) with pm.Model() as other: pm.Normal("mu", 0, 1) pm.Data("batch", np.zeros((4, 1))) @@ -285,6 +369,28 @@ def test_inference_instance_bound_to_another_model_is_rejected(): Trainer(method=elsewhere, dataloader=loader).fit(3) +def test_inference_instance_is_matched_against_the_trained_model(): + """The identity check reads the model the Trainer was given, not the context stack, + and accepts an instance built under it; the fit hands back that instance's own + Approximation, so refining or sampling the return value is refining the fit.""" + loader = DataLoader(lambda: iter(marked(10)), batch_size=4, total_size=40) + with pm.Model() as model: + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + inference = pm.ADVI(random_seed=0) + with pm.Model(): + pm.Normal("mu", 0, 1) + pm.Data("batch", np.zeros((4, 1))) + stranger = pm.ADVI() + + approx = Trainer(method=inference, dataloader=loader, model=model).fit(3) + assert approx is inference.approx + assert len(approx.hist) == 3 + with pytest.raises(ValueError, match="bound to a different model"): + Trainer(method=stranger, dataloader=loader, model=model).fit(3) + + @pytest.mark.parametrize( "declared, match", [(None, "no observed variable declares total_size"), (40, "does not match the data")], @@ -292,7 +398,7 @@ def test_inference_instance_bound_to_another_model_is_rejected(): ) def test_unusable_likelihood_scaling_warns(declared, match): """Absent or disagreeing total_size biases the posterior without failing the fit.""" - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, sample_shape=(1,), total_size=4) + loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -301,12 +407,51 @@ def test_unusable_likelihood_scaling_warns(declared, match): Trainer(method="advi", dataloader=loader).fit(3, random_seed=0) +@pytest.mark.parametrize("rows, blocks", [(4, 10), (3, 7)], ids=["N=40,batch=4", "N=21,batch=3"]) +@pytest.mark.parametrize( + "declares, match", + [ + ("absent", "no observed variable declares total_size"), + ("mismatched", "does not match the data"), + ("correct", None), + ], +) +def test_scaling_warning_fires_only_on_an_unusable_total_size(rows, blocks, declares, match): + """The scaling check warns on an absent or disagreeing total_size and stays quiet + on a correct one, for N != batch_size as well as N == batch_size. + + A false positive would be worse than the bias the check guards against: it teaches + users to ignore the warning. One shape cannot show the check reads the dataset + size rather than anything else the loader could offer. + """ + n_rows = rows * blocks + loader = DataLoader( + lambda: iter(marked(blocks, rows=rows)), + batch_size=rows, + total_size=n_rows, + ) + total_size = {"absent": None, "mismatched": n_rows + rows, "correct": n_rows}[declares] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((rows, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=total_size) + Trainer(method="advi", dataloader=loader).fit(2, random_seed=0) + scaling = [str(w.message) for w in caught if "total_size" in str(w.message)] + if match is None: + assert scaling == [] + else: + assert len(scaling) == 1 + assert match in scaling[0] + + def test_total_size_check_fires_when_fit_ends_at_pass_boundary(): """fit(n) with n exactly the batches in one pass still runs the total_size sanity check: the stream is kept one batch ahead, so stopping at the boundary does not abandon the check right before it would fire.""" data = np.zeros((40, 1)) - loader = DataLoader(chunked_factory(data, 10), batch_size=10, sample_shape=(1,), total_size=400) + loader = DataLoader(chunked_factory(data, 10), batch_size=10, total_size=400) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((10, 1))) @@ -318,9 +463,7 @@ def test_total_size_check_fires_when_fit_ends_at_pass_boundary(): def test_fit_rejects_nonpositive_n(): """fit consumes the seed batch before pm.fit could reject n itself, so a non-positive n is refused up front, before touching the stream.""" - loader = DataLoader( - lambda: iter([np.zeros((2, 1))]), batch_size=2, sample_shape=(1,), total_size=2 - ) + loader = DataLoader(lambda: iter([np.zeros((2, 1))]), batch_size=2, total_size=2) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((2, 1))) @@ -332,9 +475,7 @@ def test_fit_rejects_nonpositive_n(): def test_unknown_data_name_raises_before_consuming(): """A data_name that is not in the model raises a guided KeyError before any batch is pulled from the loader.""" - loader = DataLoader( - lambda: iter([np.zeros((4, 1))] * 3), batch_size=4, sample_shape=(1,), total_size=4 - ) + loader = DataLoader(lambda: iter([np.zeros((4, 1))] * 3), batch_size=4, total_size=4) installed = [] with pm.Model() as model: pm.Normal("mu", 0, 1) From 563dfd77088216c73e33dede29378d2b67ce99d0 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sun, 2 Aug 2026 23:36:55 -0500 Subject: [PATCH 3/9] Drop the private dataloader import and the nested stream generator Remove the cross-module `_is_positive_int` import: it is private to dataloader.py, where it validates row counts, and `n` is a step count. Inline the check instead of promoting a one-line predicate into shared code. Remove the `_stream` generator nested inside `fit`, hoisting it to a module-level `_cycle`. Epoch cycling is the loop rule that has to hold, and inside a method body it could only be exercised through a full ADVI fit; two direct tests now pin it. Remove the Notes block from the class docstring: it described a refactor of the per-step set_data as still pending, which is PR-description material rather than API documentation. Co-Authored-By: Claude --- pymc_extras/variational/trainer.py | 33 +++++++++++++----------------- tests/variational/test_trainer.py | 26 ++++++++++++++++++++++- 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py index df3b55366..58a4e83a5 100644 --- a/pymc_extras/variational/trainer.py +++ b/pymc_extras/variational/trainer.py @@ -27,11 +27,22 @@ from pymc.variational.minibatch_rv import MinibatchRandomVariable from pytensor.graph.basic import Constant -from pymc_extras.variational.dataloader import DataLoader, _is_positive_int +from pymc_extras.variational.dataloader import DataLoader __all__ = ["Trainer"] +def _cycle(loader: DataLoader) -> Iterator[np.ndarray]: + """Repeat the loader's epochs forever. A pass that yields nothing never ends.""" + while True: + empty = True + for batch in loader: + empty = False + yield batch + if empty: + raise RuntimeError("dataloader yielded no batches") + + def _warn_if_scaling_mismatches(model, total_size: int) -> None: """Warn if the model cannot rescale the minibatch likelihood to ``total_size``. @@ -94,13 +105,6 @@ class Trainer: Default keyword arguments forwarded to :func:`pymc.fit` (e.g. ``obj_optimizer``); per-call kwargs to :meth:`fit` override them. - Notes - ----- - The per-step ``set_data`` currently lives in the ``Trainer``. Once the VI - rework's ``Inference.step(batch)`` lands it moves there, at which point the - ``total_size`` rescaling can be derived from ``len(dataloader)`` and dropped - from the model body entirely. - Examples -------- .. code-block:: python @@ -149,7 +153,7 @@ def fit(self, n: int = 10_000, **kwargs): :class:`Approximation` The fitted approximation, as returned by :func:`pymc.fit`. """ - if not _is_positive_int(n): + if not isinstance(n, int) or isinstance(n, bool) or n <= 0: raise ValueError(f"n must be a positive integer (the number of fit steps), got {n!r}") loader = self.dataloader if not isinstance(loader, DataLoader): @@ -171,16 +175,7 @@ def fit(self, n: int = 10_000, **kwargs): if loader.total_size is not None: _warn_if_scaling_mismatches(model, loader.total_size) - def _stream() -> Iterator[np.ndarray]: - while True: - empty = True - for batch in loader: - empty = False - yield batch - if empty: - raise RuntimeError("dataloader yielded no batches") - - batches = _stream() + batches = _cycle(loader) model.set_data(self.data_name, next(batches)) def _advance(*_): diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index 4f6937d9b..b75f66bac 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -22,7 +22,7 @@ import pytest from pymc_extras.variational.dataloader import DataLoader -from pymc_extras.variational.trainer import Trainer +from pymc_extras.variational.trainer import Trainer, _cycle from tests.variational.dataloader_helpers import chunked_factory @@ -57,6 +57,30 @@ def streamed(loader, k): return [float(batch[0, 0]) for batch in islice(chain.from_iterable(repeat(loader)), k)] +def test_cycle_replays_the_loaders_own_epochs(): + """_cycle is exactly "repeat the loader": no batch dropped, repeated or reordered + at the epoch seam.""" + loader = DataLoader(lambda: iter(marked(3, rows=2)), batch_size=2, total_size=6) + got = [float(b[0, 0]) for b in islice(_cycle(loader), 8)] + assert got == streamed(loader, 8) + + +def test_cycle_raises_instead_of_spinning_on_an_empty_pass(): + """A source that cannot be replayed would make the cycle loop forever.""" + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + if calls["n"] == 1: + yield np.zeros((2, 1)) + + loader = DataLoader(factory, batch_size=2, total_size=2) + batches = _cycle(loader) + assert next(batches).shape == (2, 1) + with pytest.raises(RuntimeError, match="yielded no batches"): + next(batches) + + def test_trainer_end_to_end_matches_in_ram_minibatch(): """End-to-end: Trainer-driven streaming ADVI reproduces in-RAM pm.Minibatch ADVI. From 4772d644a5a847059756c362eb47b40c353535d6 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Mon, 3 Aug 2026 00:14:29 -0500 Subject: [PATCH 4/9] Pin the bool step count and the progressbar default Mutation testing found two behaviours of fit() with nothing asserting them: dropping the isinstance(n, bool) clause and dropping the progressbar setdefault both left the suite green. Co-Authored-By: Claude --- tests/variational/test_trainer.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index b75f66bac..85613fd1b 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -484,16 +484,36 @@ def test_total_size_check_fires_when_fit_ends_at_pass_boundary(): Trainer(method="advi", dataloader=loader).fit(4, random_seed=0) -def test_fit_rejects_nonpositive_n(): +@pytest.mark.parametrize("n", [0, True], ids=["zero", "bool"]) +def test_fit_rejects_nonpositive_n(n): """fit consumes the seed batch before pm.fit could reject n itself, so a - non-positive n is refused up front, before touching the stream.""" + non-positive n is refused up front, before touching the stream. ``True`` is an + ``int`` to Python, so it would otherwise pass as a one-step fit.""" loader = DataLoader(lambda: iter([np.zeros((2, 1))]), batch_size=2, total_size=2) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((2, 1))) pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) with pytest.raises(ValueError, match="positive integer"): - Trainer(method="advi", dataloader=loader).fit(0) + Trainer(method="advi", dataloader=loader).fit(n) + + +def test_progressbar_defaults_off_and_yields_to_an_explicit_setting(monkeypatch): + """pm.fit draws a progress bar by default, which a per-step streaming fit turns + into a wall of output; an explicit setting from either level still wins.""" + seen = [] + monkeypatch.setattr( + "pymc_extras.variational.trainer._fit", lambda n, **kw: seen.append(kw["progressbar"]) + ) + loader = DataLoader(lambda: iter(marked(4)), batch_size=4, total_size=16) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + trainer = Trainer(method="advi", dataloader=loader) + trainer.fit(2) + trainer.fit(2, progressbar=True) + assert seen == [False, True] def test_unknown_data_name_raises_before_consuming(): From 2ce8a1b05e94252a532199b9683c4102f9a40954 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Mon, 3 Aug 2026 03:33:03 -0500 Subject: [PATCH 5/9] Accept any integer step count in fit, and bound the empty-source fixture DataLoader validates its sizes with numbers.Integral, so fit now does the same: np.int64(4) was accepted for batch_size and refused for n. The two spin-detector fixtures now raise on a third pass, so a _cycle that loses its guard fails those tests instead of hanging them. Co-Authored-By: Claude --- pymc_extras/variational/trainer.py | 9 +++--- tests/variational/test_trainer.py | 48 ++++++++++++++++++++---------- 2 files changed, 37 insertions(+), 20 deletions(-) diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py index 58a4e83a5..9e5c67fb4 100644 --- a/pymc_extras/variational/trainer.py +++ b/pymc_extras/variational/trainer.py @@ -15,9 +15,10 @@ from __future__ import annotations +import numbers import warnings -from collections.abc import Iterator +from collections.abc import Iterable, Iterator import numpy as np @@ -32,8 +33,8 @@ __all__ = ["Trainer"] -def _cycle(loader: DataLoader) -> Iterator[np.ndarray]: - """Repeat the loader's epochs forever. A pass that yields nothing never ends.""" +def _cycle(loader: Iterable[np.ndarray]) -> Iterator[np.ndarray]: + """Repeat the loader's epochs forever, raising rather than spinning on an empty pass.""" while True: empty = True for batch in loader: @@ -153,7 +154,7 @@ def fit(self, n: int = 10_000, **kwargs): :class:`Approximation` The fitted approximation, as returned by :func:`pymc.fit`. """ - if not isinstance(n, int) or isinstance(n, bool) or n <= 0: + if not isinstance(n, numbers.Integral) or isinstance(n, bool) or n <= 0: raise ValueError(f"n must be a positive integer (the number of fit steps), got {n!r}") loader = self.dataloader if not isinstance(loader, DataLoader): diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index 85613fd1b..eb27edbf5 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -41,6 +41,24 @@ def loud(n, rows=4, spacing=1000.0): return [np.full((rows, 1), spacing * (i + 1)) for i in range(n)] +def one_epoch_then_empty(rows=4): + """A source that streams one epoch and then comes back empty. + + The third pass raises so that a ``_cycle`` which spins over the empty second one + fails the test instead of hanging it. + """ + calls = {"n": 0} + + def factory(): + calls["n"] += 1 + if calls["n"] == 1: + yield np.zeros((rows, 1)) + elif calls["n"] > 2: + raise AssertionError("_cycle restarted an exhausted source instead of raising") + + return factory + + def record_installed(model, log): """Log the marker of every batch ``set_data`` installs, in order.""" original = model.set_data @@ -67,14 +85,7 @@ def test_cycle_replays_the_loaders_own_epochs(): def test_cycle_raises_instead_of_spinning_on_an_empty_pass(): """A source that cannot be replayed would make the cycle loop forever.""" - calls = {"n": 0} - - def factory(): - calls["n"] += 1 - if calls["n"] == 1: - yield np.zeros((2, 1)) - - loader = DataLoader(factory, batch_size=2, total_size=2) + loader = DataLoader(one_epoch_then_empty(rows=2), batch_size=2, total_size=2) batches = _cycle(loader) assert next(batches).shape == (2, 1) with pytest.raises(RuntimeError, match="yielded no batches"): @@ -156,14 +167,7 @@ def test_trainer_streams_into_placeholder(): def test_trainer_raises_when_loader_cannot_restart(): """A source that streams one epoch and then comes back empty cannot be cycled; the Trainer surfaces a clear error instead of training on stale data.""" - calls = {"n": 0} - - def factory(): - calls["n"] += 1 - if calls["n"] == 1: - yield np.zeros((4, 1)) - - loader = DataLoader(factory, batch_size=4, total_size=4) + loader = DataLoader(one_epoch_then_empty(), batch_size=4, total_size=4) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) @@ -498,6 +502,18 @@ def test_fit_rejects_nonpositive_n(n): Trainer(method="advi", dataloader=loader).fit(n) +def test_fit_accepts_a_numpy_integer_step_count(): + """A numpy integer is not an ``int``, so the two modules would otherwise disagree: + the DataLoader takes one for ``batch_size`` and fit would refuse one for ``n``.""" + loader = DataLoader(lambda: iter(marked(10, rows=2)), batch_size=np.int64(2), total_size=20) + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((2, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + approx = Trainer(method="advi", dataloader=loader).fit(np.int64(3), random_seed=0) + assert len(approx.hist) == 3 + + def test_progressbar_defaults_off_and_yields_to_an_explicit_setting(monkeypatch): """pm.fit draws a progress bar by default, which a per-step streaming fit turns into a wall of output; an explicit setting from either level still wins.""" From 06b74b02048565a508dfff98c59d7dfb5d195396 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 8 Aug 2026 02:10:49 -0400 Subject: [PATCH 6/9] Follow the merged DataLoader: total_size is the property, len is the batch count The merged #698 makes len(DataLoader) the batch count, matching torch, with the dataset size N on the .total_size property. Every total_size=len(loader) in the tests silently declared N to be the batch count under the new semantics. The pass-boundary warning test pinned a loader check that the merge deleted, so it goes too. Co-Authored-By: Claude --- pymc_extras/variational/trainer.py | 19 ++++++------ tests/variational/test_trainer.py | 50 +++++++++++------------------- 2 files changed, 27 insertions(+), 42 deletions(-) diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py index 9e5c67fb4..48e9f0ebd 100644 --- a/pymc_extras/variational/trainer.py +++ b/pymc_extras/variational/trainer.py @@ -62,8 +62,8 @@ def _warn_if_scaling_mismatches(model, total_size: int) -> None: warnings.warn( "no observed variable declares total_size, so the minibatch " "log-likelihood is never rescaled and is underweighted against the prior " - "by about N / batch_size. Pass total_size=len(dataloader) to the observed " - "distribution.", + "by about N / batch_size. Pass total_size=dataloader.total_size to the " + "observed distribution.", UserWarning, stacklevel=3, ) @@ -81,8 +81,9 @@ class Trainer: Follows the design in PyMC's variational-inference rework and PyTorch Lightning: the ``Trainer`` owns the training loop, the - :class:`DataLoader` owns batching (and ``len(dataloader)`` is the dataset size - ``N``), and the model owns the math. The model exposes a ``pm.Data`` placeholder; + :class:`DataLoader` owns batching (``dataloader.total_size`` is the dataset + size ``N``; ``len(dataloader)`` is the batch count, as in torch), and the + model owns the math. The model exposes a ``pm.Data`` placeholder; the ``Trainer`` streams minibatches into it with ``model.set_data`` once per step; no user callbacks are needed. @@ -95,8 +96,8 @@ class Trainer: an instance is already bound to a model, so configure it at construction (e.g. ``ADVI(random_seed=...)``). dataloader : DataLoader - The minibatch source. ``len(dataloader)`` is ``N``; the model should pass - it to the observed distribution's ``total_size``. + The minibatch source. The model should pass ``dataloader.total_size`` to + the observed distribution's ``total_size``. model : pymc.Model, optional Defaults to the model on the context stack. data_name : str, default "batch" @@ -110,14 +111,12 @@ class Trainer: -------- .. code-block:: python - loader = DataLoader( - parquet_source("shuffled/"), batch_size=4096, sample_shape=(4,), total_size="auto" - ) + loader = DataLoader(parquet_source("shuffled/"), batch_size=4096, 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))) # placeholder 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)) + pm.Bernoulli("y", logit_p=logit, observed=batch[:, 3], total_size=loader.total_size) approx = Trainer(method="advi", dataloader=loader, data_name="batch").fit(20_000) """ diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index eb27edbf5..df753715d 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -95,7 +95,7 @@ def test_cycle_raises_instead_of_spinning_on_an_empty_pass(): def test_trainer_end_to_end_matches_in_ram_minibatch(): """End-to-end: Trainer-driven streaming ADVI reproduces in-RAM pm.Minibatch ADVI. - Exercises the whole API: a pm.Data placeholder, total_size=len(loader), and a + Exercises the whole API: a pm.Data placeholder, total_size=loader.total_size, and a Trainer that streams minibatches into the placeholder with set_data while the user writes no callbacks. Runs long enough to cycle the loader across epochs. """ @@ -135,7 +135,7 @@ def test_trainer_end_to_end_matches_in_ram_minibatch(): "o", logit_p=b[0] + b[1] * batch[:, 0] + b[2] * batch[:, 1], observed=batch[:, 2], - total_size=len(loader), + total_size=loader.total_size, ) ap = Trainer( method="advi", @@ -157,7 +157,7 @@ def test_trainer_streams_into_placeholder(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) Trainer(method="advi", dataloader=loader, data_name="batch").fit( 5, progressbar=False, random_seed=0 ) @@ -171,7 +171,7 @@ def test_trainer_raises_when_loader_cannot_restart(): with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) with pytest.raises(RuntimeError, match="yielded no batches"): Trainer(method="advi", dataloader=loader, data_name="batch").fit( 5, progressbar=False, random_seed=0 @@ -195,7 +195,7 @@ def test_trainer_appends_user_callbacks_and_streams_distinct_batches(): with pm.Model() as model: x = pm.Normal("x", 0.0, 1.0) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", x, 1.0, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", x, 1.0, observed=batch[:, 0], total_size=loader.total_size) Trainer(method="advi", dataloader=loader).fit( 5, callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))] ) @@ -211,7 +211,7 @@ def test_trainer_accepts_inference_instance(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) approx = Trainer(method=pm.ADVI(random_seed=0), dataloader=loader).fit(5) assert len(approx.hist) == 5 np.testing.assert_array_equal(model["batch"].get_value(), data) @@ -227,7 +227,7 @@ def fit_with(ctor_kwargs, fit_kwargs): with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) return Trainer(method="advi", dataloader=loader, data_name="batch", **ctor_kwargs).fit( 5, **fit_kwargs ) @@ -244,7 +244,7 @@ def test_fit_trains_one_batch_per_step(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((2, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) record_installed(model, installed) Trainer(method="advi", dataloader=loader).fit(3, random_seed=0) assert installed == [0.0, 1.0, 2.0, 3.0] @@ -275,7 +275,7 @@ def test_steps_consume_the_loaders_own_batch_sequence(n, blocks, rows): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((rows, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) record_installed(model, installed) approx = Trainer( method="advi", @@ -303,7 +303,7 @@ def test_refine_after_fit_continues_without_repeating_a_batch(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) record_installed(model, installed) inference = pm.ADVI(random_seed=0) Trainer(method=inference, dataloader=loader).fit(3) @@ -327,7 +327,7 @@ def stop_once(*_): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) record_installed(model, installed) inference = pm.ADVI(random_seed=0) Trainer(method=inference, dataloader=loader).fit(10, callbacks=[stop_once], score=False) @@ -350,7 +350,7 @@ def test_a_second_fit_reseeds_and_forgets_the_first_calls_kwargs(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) record_installed(model, installed) trainer = Trainer(method="advi", dataloader=loader, random_seed=0) unscored = trainer.fit(3, score=False) @@ -372,7 +372,7 @@ def test_user_callbacks_see_the_batch_that_produced_the_loss(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) Trainer(method="advi", dataloader=loader).fit( 5, random_seed=0, @@ -392,7 +392,7 @@ def test_inference_instance_bound_to_another_model_is_rejected(): with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) with pytest.raises(ValueError, match="bound to a different model"): Trainer(method=elsewhere, dataloader=loader).fit(3) @@ -405,7 +405,7 @@ def test_inference_instance_is_matched_against_the_trained_model(): with pm.Model() as model: mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) inference = pm.ADVI(random_seed=0) with pm.Model(): pm.Normal("mu", 0, 1) @@ -474,20 +474,6 @@ def test_scaling_warning_fires_only_on_an_unusable_total_size(rows, blocks, decl assert match in scaling[0] -def test_total_size_check_fires_when_fit_ends_at_pass_boundary(): - """fit(n) with n exactly the batches in one pass still runs the total_size - sanity check: the stream is kept one batch ahead, so stopping at the - boundary does not abandon the check right before it would fire.""" - data = np.zeros((40, 1)) - loader = DataLoader(chunked_factory(data, 10), batch_size=10, total_size=400) - with pm.Model(): - mu = pm.Normal("mu", 0, 1) - batch = pm.Data("batch", np.zeros((10, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) - with pytest.warns(UserWarning, match="disagrees with"): - Trainer(method="advi", dataloader=loader).fit(4, random_seed=0) - - @pytest.mark.parametrize("n", [0, True], ids=["zero", "bool"]) def test_fit_rejects_nonpositive_n(n): """fit consumes the seed batch before pm.fit could reject n itself, so a @@ -497,7 +483,7 @@ def test_fit_rejects_nonpositive_n(n): with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((2, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) with pytest.raises(ValueError, match="positive integer"): Trainer(method="advi", dataloader=loader).fit(n) @@ -509,7 +495,7 @@ def test_fit_accepts_a_numpy_integer_step_count(): with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((2, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) approx = Trainer(method="advi", dataloader=loader).fit(np.int64(3), random_seed=0) assert len(approx.hist) == 3 @@ -525,7 +511,7 @@ def test_progressbar_defaults_off_and_yields_to_an_explicit_setting(monkeypatch) with pm.Model(): mu = pm.Normal("mu", 0, 1) batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=len(loader)) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) trainer = Trainer(method="advi", dataloader=loader) trainer.fit(2) trainer.fit(2, progressbar=True) From ef1ded4b2a6f1d705ed00b3352d0eee871bedd4e Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 8 Aug 2026 03:15:26 -0400 Subject: [PATCH 7/9] Drop the loud guards and the design-provenance prose The DataLoader isinstance check and the data_name lookup both protected failures that announce themselves: a wrong loader dies on .total_size before anything is consumed, and a wrong data_name dies on the first set_data with a KeyError naming it, one batch in. The Inference-bound-to-another-model check stays because that failure is silent. Their tests go with them, and the class docstring loses the lineage paragraph. Co-Authored-By: Claude --- pymc_extras/variational/trainer.py | 19 +++---------------- tests/variational/test_trainer.py | 19 ------------------- 2 files changed, 3 insertions(+), 35 deletions(-) diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py index 48e9f0ebd..32098e51b 100644 --- a/pymc_extras/variational/trainer.py +++ b/pymc_extras/variational/trainer.py @@ -79,13 +79,9 @@ def _warn_if_scaling_mismatches(model, total_size: int) -> None: class Trainer: """Drive variational inference over a :class:`DataLoader` without user callbacks. - Follows the design in PyMC's variational-inference rework and PyTorch - Lightning: the ``Trainer`` owns the training loop, the - :class:`DataLoader` owns batching (``dataloader.total_size`` is the dataset - size ``N``; ``len(dataloader)`` is the batch count, as in torch), and the - model owns the math. The model exposes a ``pm.Data`` placeholder; - the ``Trainer`` streams minibatches into it with ``model.set_data`` once per - step; no user callbacks are needed. + The ``Trainer`` owns the loop, the loader owns batching, the model owns the + math: the model exposes a ``pm.Data`` placeholder and the ``Trainer`` streams + one minibatch into it per step with ``model.set_data``. Parameters ---------- @@ -156,16 +152,7 @@ def fit(self, n: int = 10_000, **kwargs): if not isinstance(n, numbers.Integral) or isinstance(n, bool) or n <= 0: raise ValueError(f"n must be a positive integer (the number of fit steps), got {n!r}") loader = self.dataloader - if not isinstance(loader, DataLoader): - raise TypeError( - f"Trainer needs a DataLoader for `dataloader`, got {type(loader).__name__}." - ) model = modelcontext(self.model) - if self.data_name not in model: - raise KeyError( - f"data_name {self.data_name!r} is not a variable in the model; it " - f"must name the pm.Data placeholder the minibatches are streamed into." - ) if isinstance(self.method, Inference) and self.method.approx.model is not model: raise ValueError( "`method` is an Inference instance bound to a different model than the " diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index df753715d..fd956c8ab 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -178,12 +178,6 @@ def test_trainer_raises_when_loader_cannot_restart(): ) -def test_trainer_rejects_non_dataloader(): - """The isinstance guard fires before any model lookup.""" - with pytest.raises(TypeError, match="DataLoader"): - Trainer(method="advi", dataloader=object()).fit(10) - - def test_trainer_appends_user_callbacks_and_streams_distinct_batches(): """User callbacks (e.g. convergence trackers) compose with the internal advance callback instead of colliding on the keyword, and the placeholder @@ -516,16 +510,3 @@ def test_progressbar_defaults_off_and_yields_to_an_explicit_setting(monkeypatch) trainer.fit(2) trainer.fit(2, progressbar=True) assert seen == [False, True] - - -def test_unknown_data_name_raises_before_consuming(): - """A data_name that is not in the model raises a guided KeyError before any - batch is pulled from the loader.""" - loader = DataLoader(lambda: iter([np.zeros((4, 1))] * 3), batch_size=4, total_size=4) - installed = [] - with pm.Model() as model: - pm.Normal("mu", 0, 1) - record_installed(model, installed) - with pytest.raises(KeyError, match=r"pm\.Data placeholder"): - Trainer(method="advi", dataloader=loader, data_name="nope").fit(2) - assert installed == [] From c8711c83e8735f3be5c6af2acbe7953931357048 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 8 Aug 2026 03:57:31 -0400 Subject: [PATCH 8/9] Pin the duck-typed loader contract the deleted isinstance guard left open The mutation audit re-added the DataLoader isinstance guard and no test failed: nothing asserted that an iterable with a total_size attribute is enough. This test trains through a plain sized iterable and kills that mutation. Co-Authored-By: Claude --- tests/variational/test_trainer.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index fd956c8ab..bb1497bfe 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -510,3 +510,21 @@ def test_progressbar_defaults_off_and_yields_to_an_explicit_setting(monkeypatch) trainer.fit(2) trainer.fit(2, progressbar=True) assert seen == [False, True] + + +def test_trainer_accepts_a_duck_typed_loader(): + """The loader contract is duck-typed -- iterable plus total_size -- so a + source that is not the concrete DataLoader class still trains.""" + + class SizedSource: + total_size = 4 + + def __iter__(self): + return iter([np.zeros((4, 1))] * 50) + + with pm.Model(): + mu = pm.Normal("mu", 0, 1) + batch = pm.Data("batch", np.zeros((4, 1))) + pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=4) + approx = Trainer(method="advi", dataloader=SizedSource()).fit(3, random_seed=0) + assert len(approx.hist) == 3 From b73687775f637320396720a4145252b188632cb0 Mon Sep 17 00:00:00 2001 From: Yicheng Yang Date: Sat, 8 Aug 2026 04:07:36 -0400 Subject: [PATCH 9/9] Say once what the Trainer spares the user, and drop two subsumed tests The module and class docstrings both read as 'user callbacks are unsupported' while fit() documents, supports, and tests callbacks=; the phrase now says what was meant, once: the user writes no hand-written streaming callbacks. test_fit_trains_one_batch_per_step and test_user_callbacks_see_the_batch_that_produced_the_loss assert strict subsets (installed only / seen only) of what the parametrized test_steps_consume_the_loaders_own_batch_sequence asserts across four n/blocks configurations; the default-optimizer sequencing path stays pinned by test_refine_after_fit_continues_without_repeating_a_batch. Co-Authored-By: Claude Fable 5 --- pymc_extras/variational/trainer.py | 4 ++-- tests/variational/test_trainer.py | 31 +----------------------------- 2 files changed, 3 insertions(+), 32 deletions(-) diff --git a/pymc_extras/variational/trainer.py b/pymc_extras/variational/trainer.py index 32098e51b..b909bba81 100644 --- a/pymc_extras/variational/trainer.py +++ b/pymc_extras/variational/trainer.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Drive variational inference over a :class:`DataLoader` with no user callbacks.""" +"""Drive variational inference over a :class:`DataLoader` without hand-written streaming callbacks.""" from __future__ import annotations @@ -77,7 +77,7 @@ def _warn_if_scaling_mismatches(model, total_size: int) -> None: class Trainer: - """Drive variational inference over a :class:`DataLoader` without user callbacks. + """Drive variational inference over a :class:`DataLoader`. The ``Trainer`` owns the loop, the loader owns batching, the model owns the math: the model exposes a ``pm.Data`` placeholder and the ``Trainer`` streams diff --git a/tests/variational/test_trainer.py b/tests/variational/test_trainer.py index bb1497bfe..afa9ae732 100644 --- a/tests/variational/test_trainer.py +++ b/tests/variational/test_trainer.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Trainer: drive variational inference over a DataLoader with no user callbacks.""" +"""Trainer: drive variational inference over a DataLoader without hand-written streaming callbacks.""" import warnings @@ -231,19 +231,6 @@ def fit_with(ctor_kwargs, fit_kwargs): np.testing.assert_array_equal(a.hist, b.hist) -def test_fit_trains_one_batch_per_step(): - """Step i trains batch i, and the step that ends the fit loads batch n for what follows.""" - loader = DataLoader(lambda: iter(marked(10, rows=2)), batch_size=2, total_size=20) - installed = [] - with pm.Model() as model: - mu = pm.Normal("mu", 0, 1) - batch = pm.Data("batch", np.zeros((2, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) - record_installed(model, installed) - Trainer(method="advi", dataloader=loader).fit(3, random_seed=0) - assert installed == [0.0, 1.0, 2.0, 3.0] - - @pytest.mark.parametrize( "n, blocks, rows", [(1, 6, 4), (4, 4, 4), (5, 3, 2), (7, 2, 1)], @@ -359,22 +346,6 @@ def test_a_second_fit_reseeds_and_forgets_the_first_calls_kwargs(): assert installed[0] != left_loaded -def test_user_callbacks_see_the_batch_that_produced_the_loss(): - """A callback reading the placeholder must see batch i on step i, not batch i+1.""" - loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4) - seen = [] - with pm.Model() as model: - mu = pm.Normal("mu", 0, 1) - batch = pm.Data("batch", np.zeros((4, 1))) - pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size) - Trainer(method="advi", dataloader=loader).fit( - 5, - random_seed=0, - callbacks=[lambda *_: seen.append(float(model["batch"].get_value()[0, 0]))], - ) - assert seen == [0.0, 1.0, 2.0, 3.0, 4.0] - - def test_inference_instance_bound_to_another_model_is_rejected(): """Streaming into one model while an Inference optimizes another is silent otherwise.""" loader = DataLoader(lambda: iter(marked(50)), batch_size=4, total_size=4)