Streaming variational inference: out-of-core DataLoader for minibatch ADVI by YichengYang-Ethan · Pull Request #698 · pymc-devs/pymc-extras

YichengYang-Ethan · 2026-06-23T01:53:25Z

Adds a streaming data layer for minibatch variational inference on data that doesn't fit in memory. Ported from pymc-devs/pymc#8325 — moving it here per the discussion on that PR (@ricardoV94's call to start in extras).

pm.Minibatch indexes an in-memory array, so peak memory is O(N). This streams minibatches from an out-of-core source into a pm.Data placeholder instead, so peak memory is set by the batch, the source chunk, and the optional shuffle buffer — independent of N.

The API mirrors torch.utils.data:

IterableDataset — a re-iterable, out-of-core source of rows (e.g. parquet_source over a directory of shards).
DataLoader — fixed-size, optionally shuffled minibatches; sized, with len(loader) == N for total_size.
shuffle_buffer — a bounded shuffle over the stream.

The unbiased-gradient rescaling reuses the existing create_minibatch_rv (the same N / batch_size as pm.Minibatch), via total_size=len(loader).

Round 2 — ran a ponytail pass over the whole diff, per review:

1361 → 1045 lines (−23%); dataloader.py 608 → 494.
The batches_seen / rows_streamed counters and the second iteration path are gone — torch's DataLoader has none and nothing read them. __iter__ is the only loop now and fires the total_size sanity check itself (that check now runs on every pass, so branch coverage went up, not down).
The duplicated Parquet schema validation is one _check_columns; _as_source returns just the factory instead of a 3-tuple.
Tests consolidated with parametrize (53 → 47, eight parametrizations); every raise/warning kept its test.
Two of the tool's suggestions were rejected on purpose — merging _rebatch with shuffle_buffer's fill loop (they differ on buffer ownership and end-of-stream flush; details in the commit message) and dropping __iter__'s one-batch lookahead (it is what makes the sanity check fire when a fit stops exactly at a pass boundary).

Verified against the pre-trim revision: the batch streams are byte-identical across plain/shuffled/Parquet/raw-array sources (same seeds), the old test suite replayed against the new code passes everywhere except the seven tests of the deleted counter API, and a streamed-ADVI run through pm.fit recovers the same posterior as in-RAM ADVI and pm.Minibatch.

Notes:

Lives at pymc_extras/variational/dataloader.py.
pyarrow is an optional dependency, imported lazily only for the Parquet source.
Tests in tests/variational/. End-to-end example: Example: out-of-core minibatch variational inference with DataLoader and Trainer pymc-examples#888 (will switch its import to pymc_extras once this lands).
The Trainer follow-up is Streaming variational inference: Trainer for minibatch ADVI #710; it consumed the deleted private stream, so it picks up a one-line adaptation there after this merges.

Ports the streaming data layer (IterableDataset, parquet_source, DataLoader, shuffle_buffer) from pymc-devs/pymc#8325. Self-contained numpy/pyarrow data layer with no pymc-internal coupling; public names mirror torch.utils.data. Tests moved alongside.

…=, sort imports

codecov-commenter · 2026-07-18T18:12:47Z

Codecov Report

❌ Patch coverage is 96.12403% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.82%. Comparing base (86fac3c) to head (e386f59).
⚠️ Report is 36 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_extras/variational/dataloader.py	96.06%	5 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #698       +/-   ##
===========================================
+ Coverage   51.60%   91.82%   +40.21%     
===========================================
  Files          73      100       +27     
  Lines        8003     9108     +1105     
===========================================
+ Hits         4130     8363     +4233     
+ Misses       3873      745     -3128

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

... and 39 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

Pretty good. I think it needs another edit. Also we should move things into a dataloader.py or something like that

zaxtax · 2026-07-18T18:07:34Z

+                f"{total_size!r}."
+            )
+
+        # Plain Python ints: create_minibatch_rv rejects np.int64 for total_size.


Drop comment

zaxtax · 2026-07-18T18:09:35Z

+            )
+        self._source_factory = source_factory
+
+        if isinstance(total_size, str):


why not just explicitly test for "auto"

zaxtax · 2026-07-18T18:09:45Z

+        seed: int | None = None,
+        sample_shape: tuple[int, ...] | None = None,
+        dtype: str = "float64",
+        total_size: int | str | None = None,


I think "auto" is a better default.

zaxtax · 2026-07-18T18:10:51Z

+            )
+        return self._total_size
+
+    def _stream_batches(self) -> Iterator[np.ndarray]:


The doc comment is a little too formally written. We can make this more clear as a comment.

zaxtax · 2026-07-18T18:11:42Z

+            yield prepared
+            batch = following
+
+    def _prepare(self, batch: np.ndarray) -> np.ndarray:


Do we need this? Does torch.DataLoader do something similar?

zaxtax · 2026-07-18T18:13:50Z

+            f"right, declare its trailing shape with DataLoader(sample_shape=...)"
+        )
+    return a
+


A lot of these factory methods look over-engineered. Maybe there is a way to tighten this logic and get rid of much of the indirection.

Move the DataLoader, IterableDataset, shuffle_buffer, and parquet_source into pymc_extras/variational/dataloader.py, with matching test files. Pure rename; no behavior change.

Following the review: - total_size now defaults to "auto" and is matched explicitly (== "auto") rather than an isinstance(str) check. - Drop the factory-of-factory indirection: _make_factory / _block_factory and the .n_rows-forwarded-onto-closures dance are replaced by a single _as_source() that normalizes any source into (new_iter, n_rows, reiterable) and reads n_rows off the original object once; shuffle_buffer no longer forwards n_rows. - Slim _prepare to preprocess + owned dtype-cast (the collate step) and drop the separate _validate (a wrong sample_shape is already caught by _promote_to_block during rebatching). - Shorten the over-formal module / _stream_batches / shuffle_buffer docstrings. - Remove the two tests that asserted the deleted n_rows forwarding; the shuffle + auto path stays covered by test_dataloader_shuffle_auto_resolves_via_n_rows. 53 DataLoader tests pass; ruff clean.

zaxtax · 2026-07-20T02:40:15Z

All these files are new. Change the copyright to 2026

zaxtax · 2026-07-20T08:46:56Z

+#   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.
+"""Out-of-core minibatching for variational inference.


Drop any mention of minibatching or scaling here

Drop the minibatching / N-over-batch-size rescaling prose from the module docstring (that belongs with the model, not the data loader) and move the usage example onto the DataLoader class docstring.

The missing-column and non-numeric-column checks were written out twice, once per shard in _ParquetDataset.__iter__ and once against the first shard in parquet_source. Both call one _check_columns helper now; both still raise, and the per-shard error still names the shard that broke. _as_source returned (new_iter, n_rows, reiterable), but the last two are one-line derivations from the source, so _auto_total_size takes the source and derives them itself and _as_source just returns the factory. Docstrings and comments lose the paragraphs that restate the module docstring or argue for a design decision; the non-obvious invariants (pyarrow silently drops unknown column names, the no-warn window for total_size, why _prepare copies) stay.

The suite had 53 tests and no parametrize, in a repo where the sibling suites lean on it (tests/inference/pathfinder/test_lbfgs.py is 6 tests in 150 lines with 4 parametrizations). Six near-identical total_size sanity tests become one four-case parametrization, and the same-contract pairs (shuffle buffer conservation, size validation, scalar sample_shape, the n_rows fast path) each become one test with ids. Four tests were strictly subsumed by a neighbour and are gone; every assertion they made still runs somewhere else. The six catch_warnings/simplefilter("error", UserWarning) blocks are redundant with filterwarnings = ["error"] in pyproject.toml, and the eight Parquet tests shared a write_parquet helper instead of repeating importorskip and write_table. Test docstrings are one line each. Branch coverage of dataloader.py is unchanged: same 7 statements and 4 partial branches missed.

Answers the review question on the counter block: torch's DataLoader has no batches_seen/rows_streamed, and nothing here read them either. Removing them leaves _stream_batches with no reason to exist separately from __iter__, so the one-batch lookahead and the total_size check move into __iter__ and the loader has a single way to iterate. The check itself is untouched, and it now runs on every pass rather than only on the accounting one, which closes the warn-once branch the tests could not reach before: 6 statements and 3 partial branches missed, down from 7 and 4. Iteration is no longer side-effect-free, so the module docstring stops claiming peak memory is one batch plus one source chunk -- measured, the lookahead holds two of each.

The total_size warning had said which value to fix; trimming it left only the diagnosis. The lazy pyarrow import lost the note saying why it is not at the top of the file, which is the kind of thing a later cleanup undoes.

The torch comparison it spelled out is already in the module docstring.

A second ponytail pass called the three shuffle integration tests one test written three times, which they were: same claim, three source shapes. One parametrization over (block factory, raw 2-D, raw 1-D scalars) makes the same assertion for all three, and a stronger one for the 2-D case, which had only checked a subset. It also proposed merging _rebatch and shuffle_buffer's loop. They only look alike: the shuffle path fills to buffer_size rather than batch_size, has to flush a final fill that never reached it, and must own the buffer it shuffles. The suite did not pin that last one -- a fill fed by a single chunk was shuffled in place with no test noticing -- so test_shuffle_buffer_does_not_mutate_source now covers both fill shapes. Also drops a variable in _auto_total_size that was set once and read once, and a caveat in shuffle_buffer's docstring that the module docstring makes.

A source is allowed to refill and re-yield one array -- the standard out-of-core reader idiom -- but _rebatch and shuffle_buffer both held such a block across further pulls and only copied at the end, so whole batches were silently replaced by later values. Iteration also ran its one-batch lookahead before preparing the current batch, aliasing it for sources that yield exact batch_size blocks. Also close two holes in the pass-size check: streaming more rows than total_size now always warns instead of getting 10% slack, and a pass too short to fill one batch says so instead of returning nothing quietly.

Sweeps one epoch over dataset sizes, batch sizes, chunkings and shuffle settings and asserts the batches are exactly floor(N/batch_size) groups of distinct source rows in source order when unshuffled -- nothing lost, duplicated or invented. Adds a second-epoch replay check, a seeded-shuffle equivalence against shuffle_buffer through the public API, and a check that a raw array, a factory, a bare re-iterable and an IterableDataset subclass all stream the same batches. Also covers three things the suite could not see before: dtype= was never asserted, preprocess_fn was never shown to run once per batch rather than per row, and the remainder a block leaves behind was never read after the source had been pulled again.

zaxtax · 2026-08-02T14:21:26Z

+        ``50 * batch_size``. A buffer as large as the dataset is a full shuffle.
+    seed : int, optional
+        Seed for the shuffle buffer (ignored when ``shuffle=False``).
+    sample_shape : tuple of int, optional


Why does DataLoader need sample_shape ? Or there any code that becomes less error-prone with it?

sample_shape existed to disambiguate a yield that could be either one sample or a block of rows: a source handing back shape (3,) might mean three scalar samples or one three-feature sample, and nothing in the data says which. But the ambiguity was self-inflicted -- it only arises because a source was allowed to yield either shape. torch avoids it by contract rather than by parameter: Dataset.__getitem__ returns one sample and the loader does the batching. Requiring blocks does the same here, and it is already what IterableDataset documents and what parquet_source has always produced. With rows on the leading axis, block.shape[1:] is the sample shape and nothing has to be declared. A raw ndarray is now handed over whole rather than iterated, since iterating one yields its rows and reintroduces exactly the ambiguity this removes. Blocks are checked for a consistent trailing shape as they are read, which turns a mid-stream shape change into a message that names both shapes rather than a numpy concatenate error later.

_prepare was a 4-line method called exactly once in __iter__; inlined. IterableDataset was an abstract base class that nothing depended on via isinstance -- _auto_total_size uses getattr(dataset, 'n_rows', None) and _as_source just calls iter(). _ParquetDataset and BlockDataset no longer subclass it.

- __len__ now returns total_size // batch_size (number of batches per epoch) - Use loader.total_size for the PyMC total_size= parameter - Simplify _auto_total_size: remove redundant 'second is first' check, inline the second iterator, tighten error messages

…batch count The merged pymc-devs#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 <noreply@anthropic.com>

… limits The num_draws + 1 proposal-pool default made importance_sampling="psis" a no-op and crashed small fits. Measured on Bayesian logistic regression against an exact full-data Laplace reference (k=5, batch 512, 200 iterations, three seeds): pool Pareto-k worst-coordinate mean gap worst sd / reference wall 1000 (None) 8.03 / 12.87 / 14.16 sd 0.93 / 0.91 / 0.92 10.5 s 1001 5.5 8.03 / 12.87 / 14.16 sd 0.93 / 0.91 / 0.92 21.7 s 2000 5.0 7.50 / 12.15 / 13.64 sd 0.80 / 0.70 / 0.76 34.6 s 4000 4.9 7.15 / 11.76 / 13.33 sd 0.74 / 0.66 / 0.71 32.8 s 8000 5.1 6.90 / 11.38 / 13.07 sd 0.71 / 0.64 / 0.68 69.8 s A pool of num_draws + 1 reproduces the unresampled numbers to every digit on every seed: importance_sampling.py:120 sets replace = (method == "psir"), so "psis" draws 1000 of 1001 without replacement, which is the pool minus one draw whatever the weights are. It also raised ValueError("n_draws_tail must be at least 5") from arviz for num_draws <= 23. The default is back to 4 * num_draws, which is fit_pathfinder's own pool (num_paths=4 x num_draws_per_path=1000 resampled to num_draws=1000), so the two APIs weight the same way and num_draws=20 fits again. It is not free and it is not clean: at this Pareto-k the 4000-draw pool has an effective sample size of 1.0-1.2 and the resampled 1000 overlap 94.5-96.1% with the top 1000 by weight, so the "reweighting" is selection, which is why the marginal sd degrades. Both facts are now in the importance_sampling docstring so the choice of None is an informed one. - Move the measured accuracy limits from the module docstring, which help() does not render, into fit_streaming_pathfinder's own docstring, and re-measure all of them. The old paragraph claimed "Pareto-k 3.4 at N=1e5" and "7.8 -> 3.9 ... at N=1e5" in the same breath, and its "sd ratio 0.8-1.1" and "fit_pathfinder 0.3-0.5" are in no measurement set on this machine. Measured today: N=1e5 gives Pareto-k 4.3-6.6 and 7.1-13.3 reference-sd; N=1.6e6 gives 22-39 and 57-101 sd; fit_pathfinder on the same N=1e5 data gives Pareto-k -0.6 to 0.4 and 0.10 sd, so the comparison is no longer "unverified". "More iterations do not help at all" was too strong: 600 instead of 200 left the selected iterate unchanged on two of three seeds and improved the third. - Plumb callbacks= through fit_streaming_pathfinder. The hook added to run_stochastic_lbfgs was unreachable through the module's only export. It also handed each callback a one-element list; pm.fit hands scores[: i + 1], so accumulate. The 1-based index was already right - _iterate_with_loss calls callback(approx, scores[: i + 1], i + s + 1) with s = 0 on a fresh fit, verified by running one. - Warn, rather than mention in passing, that a full_pass which skips the loader's preprocessing is undetectable here: measured 181 nats of logP error and no exception. - Drop the pymc_extras.variational.DataLoader reference from a public docstring; pymc-devs#722 does not depend on pymc-devs#698. - One explanation for the 1e-16 curvature floor, placed on the test it guards instead of dangling between the null-step branch and the elif. - Test the jitter (mutation M26 survived 24 tests), the callbacks path, and num_draws=20. Co-Authored-By: Claude <noreply@anthropic.com>

YichengYang-Ethan added 2 commits June 22, 2026 20:46

Fix ruff lint in streaming tests: escape regex metacharacter in match…

f1bb776

…=, sort imports

YichengYang-Ethan mentioned this pull request Jul 11, 2026

Streaming variational inference: Trainer for minibatch ADVI #710

Draft

zaxtax marked this pull request as ready for review July 18, 2026 18:06

zaxtax requested changes Jul 18, 2026

View reviewed changes

YichengYang-Ethan added 3 commits July 18, 2026 21:39

Merge remote-tracking branch 'origin/main' into streaming-dataset

f8c50e5

Rename streaming.py to dataloader.py (per review)

69f9deb

Move the DataLoader, IterableDataset, shuffle_buffer, and parquet_source into pymc_extras/variational/dataloader.py, with matching test files. Pure rename; no behavior change.

YichengYang-Ethan force-pushed the streaming-dataset branch from 6ca8eb0 to 9e2c2a1 Compare July 19, 2026 02:58

zaxtax reviewed Jul 20, 2026

View reviewed changes

YichengYang-Ethan added 7 commits July 20, 2026 06:35

Trim the module docstring and add a DataLoader example (per review)

9a11709

Drop the minibatching / N-over-batch-size rescaling prose from the module docstring (that belongs with the model, not the data loader) and move the usage example onto the DataLoader class docstring.

Keep the parts of two messages that said what to do

7febeee

The total_size warning had said which value to fix; trimming it left only the diagnosis. The lazy pyarrow import lost the note saying why it is not at the top of the file, which is the kind of thing a later cleanup undoes.

Cut the __len__ docstring down to the one thing it has to say

4d89081

The torch comparison it spelled out is already in the module docstring.

YichengYang-Ethan mentioned this pull request Jul 31, 2026

Streaming Pathfinder: minibatch L-BFGS with same-batch curvature pairs #722

Draft

YichengYang-Ethan added 2 commits August 1, 2026 16:02

zaxtax reviewed Aug 2, 2026

View reviewed changes

zaxtax self-requested a review August 6, 2026 16:10

zaxtax added 4 commits August 6, 2026 23:20

Drop all size validation — do what PyTorch does

80e69a2

Remove comments that add nothing

d2d7265

Drop meandering comment in _auto_total_size

945ef7c

zaxtax added 2 commits August 6, 2026 23:41

Trim meandering docstrings

153114e

Remove _rebatch: source must yield exact batch_size blocks

08f9d5e

zaxtax approved these changes Aug 7, 2026

View reviewed changes

zaxtax force-pushed the streaming-dataset branch from 202e686 to e386f59 Compare August 7, 2026 21:56

zaxtax merged commit 8db1880 into pymc-devs:main Aug 7, 2026
29 checks passed

Navigation Menu

Uh oh!

Streaming variational inference: out-of-core DataLoader for minibatch ADVI - #698

Streaming variational inference: out-of-core DataLoader for minibatch ADVI#698
zaxtax merged 23 commits into
pymc-devs:mainfrom
YichengYang-Ethan:streaming-dataset

YichengYang-Ethan commented Jun 23, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Jul 18, 2026 •

edited

Loading

Uh oh!

zaxtax left a comment

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax Jul 18, 2026

Uh oh!

zaxtax commented Jul 20, 2026

Uh oh!

Uh oh!

zaxtax Jul 20, 2026

Uh oh!

zaxtax Aug 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

YichengYang-Ethan commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov-commenter commented Jul 18, 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax commented Jul 20, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

YichengYang-Ethan commented Jun 23, 2026 •

edited

Loading

codecov-commenter commented Jul 18, 2026 •

edited

Loading