Streaming variational inference: Trainer for minibatch ADVI

.github/workflows/tests.yml

@@ -190,7 +190,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/variational/test_streaming_trainer.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,21 @@ 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
+   Trainer
+   shuffle_buffer
+   parquet_source
+
 .. currentmodule:: pymc
 .. autosummary::
    :toctree: generated/

pymc/variational/__init__.py

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

tests/variational/streaming_helpers.py

@@ -0,0 +1,34 @@
+#   Copyright 2024 - 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.
+"""Shared helpers for the streaming-dataset tests."""
+
+
+def chunked_factory(data, size):
+    """Return a zero-arg factory that replays ``data`` in ``size``-row chunks.
+
+    A ``DataLoader`` restarts its source once per epoch, so the source has to be
+    re-readable. This returns a *factory* (a zero-arg callable) that produces a
+    fresh generator each call, the way an out-of-core source like
+    ``parquet_source`` does; a bare generator would be one-shot and could not be
+    replayed. The
+    final chunk may hold fewer than ``size`` rows -- the loader re-batches the
+    stream to ``batch_size`` regardless -- so this also exercises the loader's
+    re-batching across uneven source blocks.
+    """
+
+    def factory():
+        for i in range(0, len(data), size):
+            yield data[i : i + size]
+
+    return factory