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

YichengYang-Ethan · 2026-06-09T03:43:44Z

Draft for mentor review (GSoC 2026, Streaming Variational Inference). Not for merge yet.

Split out per review: this PR is now the out-of-core data layer only; the Trainer follows in a stacked PR once this lands.

Adds pymc.variational.streaming, an out-of-core minibatch VI path for data that does not fit in RAM. The API matches torch.utils.data:

IterableDataset / parquet_source: a re-iterable, out-of-core source of rows. parquet_source reads one row group at a time, with the column order frozen at construction so a shard with a permuted schema cannot silently swap features.
DataLoader: batches and optionally shuffles the source. Blocks of any size are re-batched in order, and a raw array works directly: DataLoader(np.random.normal(size=(10_000, 2)), batch_size=64) from the VI overview sketch yields (64, 2) row batches (sample_shape defaults to arr.shape[1:] for an array source). Trailing rows that do not fill a final batch are dropped, like drop_last=True in torch, because the model observes a fixed-shape placeholder; with shuffle=True the dropped remainder is re-drawn each epoch, while a fixed replay order keeps dropping the same tail rows (documented, and the second open question below). The loader is sized, so total_size=len(loader) gives the N / batch_size ELBO rescaling (the same mechanism as pm.Minibatch).

With bounded source chunks the full dataset never sits in RAM at once: the model observes only a (batch_size, *sample_shape) placeholder that is overwritten with the next minibatch each step, and peak memory is bounded by the batch, the source chunks in flight, and the optional shuffle buffer (concatenating a fill briefly holds a second copy) — independent of N.

Status: 55 tests in tests/variational/test_streaming.py and test_streaming_autosize.py. The total_size log-likelihood rescaling is checked numerically (scaled by exactly N / batch_size, the same mechanism as pm.Minibatch), alongside the re-batching, shuffle-buffer, parquet_source, and total_size="auto" behavior.

Open questions for discussion:

Whether total_size can leave the model body (currently len(loader)).
Whether the loader should carry the end-of-pass remainder into the next epoch instead of dropping it, so every row participates even without shuffling.
len(loader) returns the row count N, matching the VI overview note (__len__ is total_size), not the batch count that torch returns.

Example notebook: pymc-devs/pymc-examples#888

read-the-docs-community · 2026-06-09T03:51:08Z

Documentation build overview

📚 pymc | 🛠️ Build #33447289 | 📁 Comparing 5bcab64 against latest (971800b)

🔍 Preview build

255 files changed · + 20 added · ± 230 modified · - 5 deleted

+ Added

± Modified

- Deleted

zaxtax

Leaving some comments as you prepare the draft!

zaxtax · 2026-06-10T09:30:14Z

+
+
+def test_plain_loader_rebatches_arbitrary_blocks():
+    # blocks of 3 with batch_size=4: re-batched in order; the trailing 2 rows that


Reserve comments for non-obvious code flow

Done, every test describes itself in a docstring now

zaxtax · 2026-06-10T09:30:43Z

+    batches = list(src())
+    assert all(b.shape == (10, 1) for b in batches)
+    seen = np.sort(np.concatenate([b.ravel() for b in batches]))
+    # 140 rows, batch 10 -> 14 full batches, nothing dropped (140 % 10 == 0)


Remove this comment

zaxtax · 2026-06-10T09:31:52Z

+    np.testing.assert_array_equal(ds.as_tensor().get_value(), np.full((4, 1), 3.0))
+
+
+def test_shuffle_buffer_conserves_rows_non_dividing():


Use docstrings not comments for describing the test

zaxtax · 2026-06-10T09:32:18Z

+
+
+def test_total_size_rescales_logp_like_minibatch():
+    # observed=buf[:, k] + total_size=N must scale the observed log-likelihood by


Unclear comment

zaxtax · 2026-06-10T09:32:54Z

+        DataLoader(_chunks(data, 4), batch_size=True, sample_shape=(1,), total_size=8)
+
+
+def test_shuffle_buffer_reshuffles_across_epochs():


This comment would be better done with a more descriptive test name

zaxtax · 2026-06-10T09:33:51Z

+            np.zeros((batch_size, *self._sample_shape), dtype=dtype), name=name
+        )
+
+    # ----- read-only state ---------------------------------------------------


Remove all these banners

zaxtax · 2026-06-10T09:34:46Z

+        """
+        return self._shared
+
+    def advance(self) -> None:


Is advance used in any public API here or in pytorch?

No, torch's DataLoader has no advance and there's no public one here; I removed that path the only advance left is a private _advance closure inside Trainer.fit that drives set_data

zaxtax · 2026-06-10T09:35:13Z

+
+    Notes
+    -----
+    This is the *starting point* Rob suggested: the per-step ``set_data`` logic


I hope this note won't be here once it's draft ready ;)

zaxtax · 2026-06-10T09:36:13Z

+
+        Returns whatever :func:`pymc.fit` returns for the chosen method.
+        """
+        from pymc.model import modelcontext


I would do import as instead of an internal import here.

zaxtax · 2026-06-10T09:36:54Z

+
+
+def test_dataloader_shuffle_true_yields_full_batches():
+    # shuffle=True wraps the source in a bounded shuffle_buffer internally; batches


Should be a docstring not a comment. Please apply this to the rest of the tests

zaxtax · 2026-06-14T16:40:30Z

+)
+
+
+def _chunks(data, size):


Why is this here?

fresh iterator each call (a bare generator would be one-shot). It yields the data a chunk of rows at a time, the way an out-of-core source like I can fold into a conftest if you'd prefer.

codecov · 2026-06-14T16:50:45Z

Codecov Report

❌ Patch coverage is 81.70732% with 45 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.62%. Comparing base (c032457) to head (5bcab64).

Files with missing lines	Patch %	Lines
pymc/variational/streaming.py	81.63%	45 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8325      +/-   ##
==========================================
- Coverage   91.76%   91.62%   -0.14%     
==========================================
  Files         128      129       +1     
  Lines       20653    20899     +246     
==========================================
+ Hits        18952    19149     +197     
- Misses       1701     1750      +49

Files with missing lines	Coverage Δ
pymc/variational/__init__.py	`100.00% <100.00%> (ø)`
pymc/variational/streaming.py	`81.63% <81.63%> (ø)`

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

fonnesbeck · 2026-06-19T15:47:48Z

@ricardoV94 Do we want this here or in extras?

ricardoV94 · 2026-06-21T20:56:26Z

If it can go in extras, let's go there first

zaxtax · 2026-07-05T07:57:29Z

@YichengYang-Ethan can you rebase this against main and mark it ready for review?

pm.Minibatch random-indexes a fully-resident array (peak memory O(N)). StreamingDataset feeds minibatches from an arbitrary source into a small pytensor.shared buffer (peak memory O(batch_size)), reusing the existing total_size / create_minibatch_rv rescaling unchanged. Adds a shuffle_buffer helper and an equivalence test (streaming ADVI == in-RAM pm.Minibatch ADVI).

Close three silent-corruption holes found in a 5-lens review: - reject total_size <= 0 in __init__: 0 is falsy and skips the N/batch_size rescaling entirely (posterior collapses to prior); negative flips the data log-likelihood's sign via get_scaling. - shuffle_buffer now accumulates max(buffer_size, batch_size) rows before emitting, so buffer_size < batch_size no longer silently discards the whole stream; also validate buffer_size/batch_size as positive ints. - positive-int checks use numbers.Integral (accept numpy ints, reject bool). +5 regression tests; existing 10 unchanged and passing.

A seeded shuffle_buffer rebuilt its RNG from the same seed on every factory call, so under cycle=True every epoch replayed one fixed permutation -- which weakens the very mixing the buffer exists to provide and compounds the block-shuffle bias on ordered data. Derive a fresh per-epoch sub-stream from a SeedSequence so the order differs across epochs while staying reproducible for a given seed. +2 tests.

Cuts the "user must pass total_size" burden (open question pymc-devs#1 for the design review): - total_size="auto" resolves N from a source's .n_rows (cheap -- e.g. Parquet footer metadata via the new parquet_source) else one counting pass over a finite, re-readable source. One-shot / infinite sources still pass total_size explicitly (and are rejected with a clear error under "auto"). - a free sanity check using the existing rows_streamed counter: at the first epoch boundary, warn if total_size grossly disagrees with the rows actually streamed in one pass (catches a wrong-but-positive total_size). - parquet_source(directory): a finite, re-readable source carrying .n_rows read from Parquet metadata (no data scan). +7 tests; the existing 17 are unchanged and still pass.

…tion An adversarial re-review surfaced edge cases the first hardening pass missed: - total_size / batch_size: numpy integers were accepted but stored unchanged, so a stored np.int64 reached create_minibatch_rv and raised "Invalid type for total_size". Normalize to Python int at construction. - _make_factory: a zero-arg factory returning a non-iterator iterable (e.g. a list) crashed in __next__ ("'list' object is not an iterator"); wrap in iter(). - total_size="auto": a factory that returns the same one-shot iterator each call now raises, instead of leaving the first advance() empty. - fit_callback: seeds the buffer by default. PyMC runs callbacks after each step, so an unseeded first step trained on the zero-initialized placeholder. - _validate: a 0-D batch now raises a clear ValueError instead of IndexError. Adds 7 regression tests (31 total).

…ze="auto" shuffle_buffer now propagates a known .n_rows (e.g. parquet_source's, read from Parquet metadata) to its wrapped factory, so the common composition StreamingDataset(shuffle_buffer(parquet_source(dir)), total_size="auto") resolves N for free instead of doing a full counting pass over the data. The only discrepancy is the single dropped trailing partial batch (< batch_size rows), which is within the auto-size sanity tolerance. Adds 2 regression tests (33 total).

…iner Design-review feedback from Rob (mentor): the streaming API should mirror torch.utils.data so the mental model transfers, and the user-facing callback should go away in favour of a Lightning-style trainer. - IterableDataset: re-iterable out-of-core source base (parquet_source now returns one); carries an optional .n_rows for total_size="auto". - DataLoader: the former StreamingDataset, renamed; gains PyTorch-style shuffle=/buffer_size=/seed= (wraps shuffle_buffer internally). Still owns the fixed pytensor.shared buffer the model observes; advance()/as_tensor() kept. - Trainer: Trainer(method="advi").fit(model, loader, n) drives VI with NO user-facing callbacks -- it seeds the buffer and advances it each step internally. The per-step advance is wired into pm.fit privately. All hardening preserved (int normalization, total_size guards + "auto", shuffle row-conservation + per-epoch reshuffle, copy-before-borrow, validation). shuffle_buffer/parquet_source stay public. 36 tests pass (1 skipped: pyarrow). total_size still appears in the model (total_size=loader.total_size); removing it is an open design question for Rob -- see notes. It is compiled into the logp graph at register_rv time (MinibatchRandomVariable Op), so fit-time injection needs either Trainer graph surgery or a dims-based rule in core.

Follows jessegrabowski/pymc VI_Overview.ipynb (the VI rework Rob/Jesse are building) instead of my ad-hoc shapes: - DataLoader.__len__ == total_size N (sized like a PyTorch DataLoader), and __iter__ yields the validated minibatch stream. This is the answer to Rob's open question: total_size leaves the model and becomes len(loader). - Trainer takes (method=, dataloader=, model=, data_name=) and fit(n); it streams each minibatch into the model's pm.Data placeholder via model.set_data, so the model is fully decoupled from the loader and the user writes no callbacks. - Model idiom is now pm.Data("batch", placeholder) + total_size=len(loader), matching the blueprint; verified end-to-end (recovers in-RAM pm.Minibatch ADVI). - Kept the as_tensor()/advance() shared-buffer path as a documented advanced escape hatch; dropped the now-unused _seed_buffer/_advance_callback. 38 tests pass (1 skipped: pyarrow). Open for Rob: spelling DataLoader (PyTorch, per his "match PyTorch") vs Dataloader (Jesse's draft); method-as-string until the ADVI(Inference).step rework lands.

- Trainer's stream now updates batches_seen/rows_streamed and runs the one-shot total_size sanity check at each epoch boundary (previously dead on the Trainer path; __iter__ stays side-effect-free). - total_size="auto" with shuffle=True counts the unshuffled source, fixing an undercount of up to batch_size-1 rows. - Trainer default data_name "data" -> "batch" to match the examples/tests. - Clarify len(loader)==N (rows, not batches) in docstrings; raise a clear error when a cycled source restarts empty. - Register the streaming API in docs/source/api/vi.rst. - Add regression tests for the auto-size shuffle count and Trainer counters.

The non-shuffle path previously required the source to yield exact batch_size blocks and raised on anything else, while the docstrings promised re-batching. Now both paths re-batch: blocks of any size are sliced in order with remainders carried across blocks, and a raw array (or any single-sample stream) is accepted directly, so the VI-rework sketch usage Dataloader(<array>, batch_size=...) works as written. Trailing rows that do not fill a final batch are dropped, like drop_last=True in torch, since the model observes a fixed-shape placeholder. Also: total_size="auto" counts a single-sample stream as rows rather than flattened elements; Trainer.fit(callbacks=...) appends user callbacks after the internal advance instead of raising a duplicate keyword error.

- Drop the shared-buffer path (as_tensor/advance and the cycle/name parameters): neither exists in torch.utils.data and the Trainer never used it. Manual stepping stays available through plain iteration plus set_data. - Move modelcontext/fit imports to module level. - Replace test comments with docstrings, drop redundant comments and section banners, and rename the reshuffle test descriptively.

shuffle_buffer concatenates yields along the leading axis, so a raw array source under shuffle=True had its rows flattened (2-D) or crashed on shape[0] (scalars). Promote single samples to one-row blocks before the shuffle wrap, with the same helper the re-batcher uses. Also tighten a few docstring claims: the parquet dtype follows the file columns, and the shuffle buffer bound is stated as rows held.

DataLoader infers sample_shape from a raw array source, so DataLoader(arr, batch_size=...) batches rows instead of silently flattening them to scalars. The total_size check no longer warns on an exact N when drop-last truncates the final batch, and its advice covers a wrong source n_rows. Trainer.fit routes all kwargs through one merge so constructor defaults like random_seed work as documented, accepts an Inference instance, and rejects an unknown data_name before consuming a batch. parquet_source validates columns against the schema up front. The shuffle_buffer docstring states the true buffer bound.

- Trainer.fit(n) consumes exactly n batches: the advance after the final step is skipped, so a finite source is not over-consumed - the total_size sanity check counts the pass that completed instead of the cumulative row counter, which inflated across partial streams - parquet_source freezes the column order at construction and reads one row group at a time, so a permuted shard schema cannot silently swap features and peak read memory is a row group, not a file - warn at construction when a fixed-order loader would drop the same non-divisible tail every pass - total_size='auto' probes that the factory can actually be re-read, catching factories that close over an already-consumed iterator - document the shuffle-buffer transient concatenation copy and the full-buffer case

…diagnostics - the internal advance skips only fit's own final step, so Inference.refine on a method instance keeps streaming instead of silently retraining on the last batch - keep the rebatcher one batch ahead in the accounting stream, so the total_size sanity check still fires when fit(n) stops exactly at the pass boundary - drop the fixed-order divisibility warning: it false-alarmed on the module's own pre-shuffled-on-disk example and on manual shuffle_buffer wrapping; the drop-last caveat lives in the docs instead - validate n in Trainer.fit (fit(0) consumed the seed batch; fit(-1) failed deep inside PyTensor) - normalize shuffle_buffer's factory output with iter(), which a re-iterable-returning factory would otherwise restart every fill - parquet_source rejects non-numeric columns at construction and names the shard when a later file is missing a frozen column - name the sample_shape remedy in the block-shape error; spell behavior consistently

The class summary still claimed the full dataset never enters memory in the absolute; match the module docstring's bounded-source-chunks framing and fix a double space.

- _ParquetDataset checks each shard's column types, so a later shard whose column turned non-numeric is named instead of failing as an opaque float cast downstream (parquet_source only saw the first shard) - the fit docstring no longer says 'exactly n consumed'; it feeds exactly n batches to the model, but the one-batch lookahead can read a re-readable source one batch further - the refine test now uses distinct batch markers and pins the honest resume-not-reseed behavior (its first step reuses fit's last batch) instead of only checking counters on all-ones data

Keep this PR to the dataset/loader layer (IterableDataset, DataLoader, shuffle_buffer, parquet_source); the Trainer and its tests move to a stacked follow-up PR. Fold the re-readable chunked-source factory the loader tests share into tests/variational/streaming_helpers.py, which doubles as a place to explain why a re-readable factory (not a one-shot generator) is needed.

YichengYang-Ethan · 2026-07-05T14:08:35Z

@YichengYang-Ethan can you rebase this against main and mark it ready for review?

Done, rebased against main and marked it ready.

ricardoV94 · 2026-07-06T06:58:53Z

isn't this going to extras?

zaxtax · 2026-07-06T07:01:35Z

oh wait yes

…

On Mon, 6 Jul 2026, 08:59 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc#8325) <#8325 (comment)> isn't this going to extras? — Reply to this email directly, view it on GitHub <#8325?email_source=notifications&email_token=AAACCUMNJGKXJ6IJTGNBSFD5DNE4HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYHE4TOMRYGY32M4TFMFZW63VHMNXW23LFNZ2KKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-4889972867>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUOIHBOMXGLMF3VFKO35DNE4HAVCNFSNUABCKJSXA33TNF2G64TZHMYTSMRZGA2DWSLTON2WKOZUGYYTQNJZGY2DIMVBOYBA> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/AAACCUN6FJMMWN4S47RDERT5DNE4HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYHE4TOMRYGY32M4TFMFZW63VHMNXW23LFNZ2KKZLWMVXHJKTGN5XXIZLSL5UW64Y> and Android <https://github.com/notifications/mobile/android/AAACCUP7MAHABQBFCSQJEM35DNE4HA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBYHE4TOMRYGY32M4TFMFZW63VHMNXW23LFNZ2KKZLWMVXHJLTGN5XXIZLSL5QW4ZDSN5UWI>. Download it today! You are receiving this because you commented.Message ID: ***@***.***>

YichengYang-Ethan · 2026-07-12T16:09:59Z

Ported to pymc-extras as #698 following the extras-first discussion — closing the core version to keep review in one place. Thanks!

… ADVI (#698) * Add out-of-core DataLoader for minibatch variational inference 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. * Fix ruff lint in streaming tests: escape regex metacharacter in match=, sort imports * Rename streaming.py to dataloader.py (per review) Move the DataLoader, IterableDataset, shuffle_buffer, and parquet_source into pymc_extras/variational/dataloader.py, with matching test files. Pure rename; no behavior change. * Address DataLoader review: auto default, less indirection, casual docs 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. * Set copyright year to 2026 on the new files (per review) * Trim the module docstring and add a DataLoader example (per review) 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. * Share one Parquet schema check and flatten the source normalizer 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. * Consolidate the DataLoader tests with parametrize 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. * Drop the batch counters and give the loader one iteration path 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. * Keep the parts of two messages that said what to do 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 The torch comparison it spelled out is already in the module docstring. * Fold the three shuffle=True tests into one, and pin the copy 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. * Copy source blocks before the source is pulled again 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. * Pin the DataLoader's conservation and equivalence invariants 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. * Take the sample shape from the blocks instead of a constructor argument 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. * Inline _prepare and drop IterableDataset _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. * Drop all size validation — do what PyTorch does * Remove comments that add nothing * Drop meandering comment in _auto_total_size * Trim meandering docstrings * Remove _rebatch: source must yield exact batch_size blocks * Match torch: len(DataLoader) returns batch count, not N - __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 --------- Co-authored-by: Rob Zinkov <rob@zinkov.com>

YichengYang-Ethan mentioned this pull request Jun 9, 2026

Example: out-of-core minibatch variational inference with DataLoader and Trainer pymc-devs/pymc-examples#888

Closed

zaxtax reviewed Jun 10, 2026

View reviewed changes

YichengYang-Ethan force-pushed the streaming-dataset-draft branch from aa560c8 to 13a6a05 Compare June 11, 2026 04:53

zaxtax reviewed Jun 14, 2026

View reviewed changes

YichengYang-Ethan changed the title ~~Streaming variational inference: out-of-core DataLoader + Trainer for minibatch ADVI~~ Streaming variational inference: out-of-core DataLoader for minibatch ADVI Jun 16, 2026

YichengYang-Ethan mentioned this pull request Jun 16, 2026

Streaming variational inference: Trainer for minibatch ADVI #8333

Closed

YichengYang-Ethan mentioned this pull request Jun 23, 2026

Streaming variational inference: out-of-core DataLoader for minibatch ADVI pymc-devs/pymc-extras#698

Merged

YichengYang-Ethan added 18 commits July 5, 2026 09:02

Register streaming tests in CI matrix (check_all_tests_are_covered)

8fef5a7

Fix mypy error comparing Integral with int

9cc789a

Match docstring punctuation to the rest of the codebase

12b0d7c

Follow numpydoc section conventions

2eb0d6f

Declare __all__ like the neighboring modules

72034e9

Add docstrings to the last three tests

67d438a

Plainer docstring wording

fd6931d

YichengYang-Ethan added 9 commits July 5, 2026 09:02

State the memory bound precisely

0596671

Apply formatter

fcdb80c

Align the DataLoader summary with the module's bounded-chunks wording

6167c30

The class summary still claimed the full dataset never enters memory in the absolute; match the module docstring's bounded-source-chunks framing and fix a double space.

YichengYang-Ethan force-pushed the streaming-dataset-draft branch from 796306d to 5bcab64 Compare July 5, 2026 14:02

YichengYang-Ethan marked this pull request as ready for review July 5, 2026 14:06

YichengYang-Ethan mentioned this pull request Jul 11, 2026

Streaming variational inference: Trainer for minibatch ADVI pymc-devs/pymc-extras#710

Open

YichengYang-Ethan closed this Jul 12, 2026



		def test_plain_loader_rebatches_arbitrary_blocks():
		# blocks of 3 with batch_size=4: re-batched in order; the trailing 2 rows that

		np.testing.assert_array_equal(ds.as_tensor().get_value(), np.full((4, 1), 3.0))


		def test_shuffle_buffer_conserves_rows_non_dividing():



		def test_total_size_rescales_logp_like_minibatch():
		# observed=buf[:, k] + total_size=N must scale the observed log-likelihood by

		DataLoader(_chunks(data, 4), batch_size=True, sample_shape=(1,), total_size=8)


		def test_shuffle_buffer_reshuffles_across_epochs():



		def test_dataloader_shuffle_true_yields_full_batches():
		# shuffle=True wraps the source in a bounded shuffle_buffer internally; batches

Uh oh!

Conversation

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

Uh oh!

Uh oh!

read-the-docs-community Bot commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

fonnesbeck commented Jun 19, 2026

Uh oh!

ricardoV94 commented Jun 21, 2026

Uh oh!

zaxtax commented Jul 5, 2026

Uh oh!

YichengYang-Ethan commented Jul 5, 2026

Uh oh!

ricardoV94 commented Jul 6, 2026

Uh oh!

zaxtax commented Jul 6, 2026 via email

Uh oh!

YichengYang-Ethan commented Jul 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

YichengYang-Ethan commented Jun 9, 2026 •

edited

Loading

read-the-docs-community Bot commented Jun 9, 2026 •

edited

Loading

codecov Bot commented Jun 14, 2026 •

edited

Loading