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

.github/workflows/tests.yml

@@ -192,7 +192,7 @@ jobs:
         linker: [cvm, numba]
         python-version: ["3.12"]
         test-subset:
-          - tests/variational/test_approximations.py tests/variational/test_callbacks.py tests/variational/test_inference.py tests/variational/test_opvi.py tests/test_initial_point.py
+          - tests/variational/test_approximations.py tests/variational/test_callbacks.py tests/variational/test_inference.py tests/variational/test_opvi.py tests/variational/test_streaming.py tests/variational/test_streaming_autosize.py tests/test_initial_point.py
           - tests/model/test_core.py tests/sampling/test_mcmc.py
           - tests/gp/test_cov.py tests/gp/test_gp.py tests/gp/test_mean.py tests/gp/test_util.py tests/ode/test_ode.py tests/ode/test_utils.py tests/smc/test_smc.py tests/sampling/test_parallel.py
           - tests/step_methods/test_metropolis.py tests/step_methods/test_slicer.py tests/step_methods/hmc/test_nuts.py tests/step_methods/test_compound.py tests/step_methods/hmc/test_hmc.py tests/step_methods/test_state.py

docs/source/api/vi.rst

@@ -68,6 +68,20 @@ Special
 
    Stein
 
+Streaming
+---------
+Out-of-core minibatching for variational inference on datasets that do not fit in
+memory (see :mod:`pymc.variational.streaming`).
+
+.. currentmodule:: pymc.variational
+.. autosummary::
+   :toctree: generated/
+
+   DataLoader
+   IterableDataset
+   shuffle_buffer
+   parquet_source
+
 .. currentmodule:: pymc
 .. autosummary::
    :toctree: generated/

pymc/variational/__init__.py

@@ -44,6 +44,12 @@
 
 # special
 from pymc.variational.stein import Stein
+from pymc.variational.streaming import (
+    DataLoader,
+    IterableDataset,
+    parquet_source,
+    shuffle_buffer,
+)
 from pymc.variational.updates import (
     adadelta,
     adagrad,
@@ -64,10 +70,12 @@
     "ADVI",
     "ASVGD",
     "SVGD",
+    "DataLoader",
     "Empirical",
     "FullRank",
     "FullRankADVI",
     "Group",
+    "IterableDataset",
     "MeanField",
     "adadelta",
     "adagrad",
@@ -80,8 +88,10 @@
     "momentum",
     "nesterov_momentum",
     "norm_constraint",
+    "parquet_source",
     "rmsprop",
     "sample_approx",
     "sgd",
+    "shuffle_buffer",
     "total_norm_constraint",
 )