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

Closed

YichengYang-Ethan wants to merge 27 commits into

pymc-devs:mainfrom

YichengYang-Ethan:streaming-dataset-draft

Commits on Jul 5, 2026

feat(variational): StreamingDataset for out-of-core minibatch VI
YichengYang-Ethan
committed
6b05063
harden StreamingDataset validation (deep-review fixes)
YichengYang-Ethan
committed
f336456
strengthen shuffle_buffer: reshuffle each epoch
YichengYang-Ethan
committed
2c43cfc
prototype: total_size="auto" + rows_streamed sanity warning
YichengYang-Ethan
committed
5ccb6d6
fix(streaming): normalize int sizes; harden factory, callback, validation
YichengYang-Ethan
committed
d6ed159
feat(streaming): forward shuffle_buffer's source .n_rows for total_size="auto"
YichengYang-Ethan
committed
9b092c5
refactor(streaming): adopt PyTorch-style Dataset/DataLoader + add Trainer
YichengYang-Ethan
committed
0294f0c
refactor(streaming): align Trainer/DataLoader with VI-rework blueprint
YichengYang-Ethan
committed
318886d
Refine streaming VI after mentor design review
YichengYang-Ethan
committed
626a2e4
Register streaming tests in CI matrix (check_all_tests_are_covered)
YichengYang-Ethan
committed
8fef5a7
Re-batch arbitrary block sizes in the plain DataLoader path
YichengYang-Ethan
committed
bcea674
Address review comments
YichengYang-Ethan
committed
52c2322
Fix mypy error comparing Integral with int
YichengYang-Ethan
committed
9cc789a
Match docstring punctuation to the rest of the codebase
YichengYang-Ethan
committed
12b0d7c
Follow numpydoc section conventions
YichengYang-Ethan
committed
2eb0d6f
Declare __all__ like the neighboring modules
YichengYang-Ethan
committed
72034e9
Add docstrings to the last three tests
YichengYang-Ethan
committed
67d438a
Plainer docstring wording
YichengYang-Ethan
committed
fd6931d
State the memory bound precisely
YichengYang-Ethan
committed
0596671
Promote single samples before the shuffle buffer
YichengYang-Ethan
committed
20b6970
Apply formatter
YichengYang-Ethan
committed
fcdb80c
Fix loader and trainer edge cases; tighten docstring claims
YichengYang-Ethan
committed
f6d4fee
Fix batch accounting and parquet edge cases from review
YichengYang-Ethan
committed
c0b39c8
Address fourth review round: refine support, boundary check, parquet diagnostics
YichengYang-Ethan
committed
8921248
Align the DataLoader summary with the module's bounded-chunks wording
YichengYang-Ethan
committed
6167c30
Tighten parquet type checks, fit/refine wording from review
YichengYang-Ethan
committed
2fb40b3
Split the Trainer into a follow-up PR; share the loader test helper
YichengYang-Ethan
committed
5bcab64

Navigation Menu

Uh oh!

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

Streaming variational inference: out-of-core DataLoader for minibatch ADVI#8325
YichengYang-Ethan wants to merge 27 commits into
pymc-devs:mainfrom
YichengYang-Ethan:streaming-dataset-draft

Commits on Jul 5, 2026

feat(variational): StreamingDataset for out-of-core minibatch VI

harden StreamingDataset validation (deep-review fixes)

strengthen shuffle_buffer: reshuffle each epoch

prototype: total_size="auto" + rows_streamed sanity warning

fix(streaming): normalize int sizes; harden factory, callback, validation

feat(streaming): forward shuffle_buffer's source .n_rows for total_size="auto"

refactor(streaming): adopt PyTorch-style Dataset/DataLoader + add Trainer

refactor(streaming): align Trainer/DataLoader with VI-rework blueprint

Refine streaming VI after mentor design review

Register streaming tests in CI matrix (check_all_tests_are_covered)

Re-batch arbitrary block sizes in the plain DataLoader path

Address review comments

Fix mypy error comparing Integral with int

Match docstring punctuation to the rest of the codebase

Follow numpydoc section conventions

Declare all like the neighboring modules

Add docstrings to the last three tests

Plainer docstring wording

State the memory bound precisely

Promote single samples before the shuffle buffer

Apply formatter

Fix loader and trainer edge cases; tighten docstring claims

Fix batch accounting and parquet edge cases from review

Address fourth review round: refine support, boundary check, parquet diagnostics

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

Tighten parquet type checks, fit/refine wording from review

Split the Trainer into a follow-up PR; share the loader test helper