diff --git a/pymc_extras/variational/__init__.py b/pymc_extras/variational/__init__.py
new file mode 100644
index 000000000..782872e7c
--- /dev/null
+++ b/pymc_extras/variational/__init__.py
@@ -0,0 +1,24 @@
+#   Copyright 2026 - 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.
+from pymc_extras.variational.dataloader import (
+    DataLoader,
+    parquet_source,
+    shuffle_buffer,
+)
+
+__all__ = [
+    "DataLoader",
+    "parquet_source",
+    "shuffle_buffer",
+]
diff --git a/pymc_extras/variational/dataloader.py b/pymc_extras/variational/dataloader.py
new file mode 100644
index 000000000..6e22b9964
--- /dev/null
+++ b/pymc_extras/variational/dataloader.py
@@ -0,0 +1,308 @@
+#   Copyright 2026 - 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.
+"""Stream out-of-core data into a PyMC model, one batch at a time.
+
+The API mirrors ``torch.utils.data``: a re-iterable source of rows is turned into
+fixed-size, optionally shuffled batches by a :class:`DataLoader`. A source yields
+blocks of rows: the leading axis is the rows, so ``block.shape[1:]`` is one sample.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+import warnings
+
+from collections.abc import Callable, Iterable, Iterator
+
+import numpy as np
+
+__all__ = ["DataLoader", "parquet_source", "shuffle_buffer"]
+
+
+def _as_source(
+    dataset: Iterable[np.ndarray] | Callable[[], Iterator[np.ndarray]],
+) -> Callable[[], Iterator[np.ndarray]]:
+    """Normalize any accepted source into a zero-arg factory returning a fresh iterator."""
+    if isinstance(dataset, Iterator):
+        used = {"done": False}
+
+        def new_iter() -> Iterator[np.ndarray]:
+            if used["done"]:
+                raise RuntimeError(
+                    "source is a bare iterator and was already consumed; the loader "
+                    "restarts the stream each epoch, so pass a zero-arg factory or a "
+                    "re-iterable instead"
+                )
+            used["done"] = True
+            return dataset
+
+        return new_iter
+
+    if isinstance(dataset, np.ndarray):
+        # Iterating an array yields its rows one at a time, which loses the
+        # distinction between a row and a block; hand it over whole instead.
+        return lambda: iter((dataset,))
+
+    make = dataset if callable(dataset) else (lambda: dataset)
+    return lambda: iter(make())
+
+
+def _auto_total_size(
+    dataset: Iterable[np.ndarray] | Callable[[], Iterator[np.ndarray]],
+    new_iter: Callable[[], Iterator[np.ndarray]],
+) -> int:
+    """Resolve ``total_size="auto"``: trust a source ``.n_rows``, else count once."""
+    n_rows = getattr(dataset, "n_rows", None)
+    if n_rows is not None:
+        return int(n_rows)
+    if isinstance(dataset, Iterator):
+        raise ValueError(
+            "total_size='auto' needs a re-readable source (a zero-arg factory or an "
+            "iterable), not a one-shot iterator; pass total_size=N explicitly instead."
+        )
+    warnings.warn(
+        "total_size='auto' is doing a full counting pass over the source; for a cheap "
+        "path use a source exposing .n_rows (e.g. parquet_source, from Parquet metadata).",
+        UserWarning,
+        stacklevel=3,
+    )
+    first = new_iter()
+    count = 0
+    for chunk in first:
+        count += int(np.asarray(chunk).shape[0])
+    if count <= 0:
+        raise ValueError("total_size='auto' counted 0 rows (empty source).")
+    if next(new_iter(), None) is None:
+        raise ValueError(
+            "total_size='auto' counted rows but the source is not re-readable "
+            "(it returns a one-shot iterator, or closes over an already-consumed one); "
+            "pass a source that makes a fresh iterator each epoch, "
+            "or total_size=N explicitly."
+        )
+    return count
+
+
+def shuffle_buffer(
+    chunk_source: Callable[[], Iterator[np.ndarray]],
+    *,
+    buffer_size: int,
+    batch_size: int,
+    seed: int | None = None,
+) -> Callable[[], Iterator[np.ndarray]]:
+    """Wrap a block source into a shuffled, fixed-size batch source.
+
+    Fills a buffer of at least ``buffer_size`` rows, shuffles, and yields
+    ``batch_size`` slices. Each epoch draws a fresh permutation from ``seed``.
+    """
+    seed_seq = np.random.SeedSequence(seed)
+
+    def factory() -> Iterator[np.ndarray]:
+        rng = np.random.default_rng(seed_seq.spawn(1)[0])
+        # chunk_source() may be any re-iterable; normalize to one iterator so each
+        # fill continues the stream instead of restarting it.
+        it = iter(chunk_source())
+        carry: np.ndarray | None = None
+        exhausted = False
+        # Accumulate at least one batch even if buffer_size < batch_size, else the
+        # guard below would discard the whole stream.
+        target = max(buffer_size, batch_size)
+        while not exhausted:
+            bufs: list[np.ndarray] = []
+            have = 0
+            if carry is not None:
+                bufs.append(carry)
+                have += carry.shape[0]
+                carry = None
+            for arr in it:
+                a = np.array(arr)
+                bufs.append(a)
+                have += a.shape[0]
+                if have >= target:
+                    break
+            else:
+                exhausted = True
+            if have < batch_size:
+                return
+            buf = np.concatenate(bufs, axis=0)
+            rng.shuffle(buf)
+            n_full = buf.shape[0] // batch_size
+            for i in range(n_full):
+                yield buf[i * batch_size : (i + 1) * batch_size]
+            rem = buf.shape[0] - n_full * batch_size
+            carry = buf[n_full * batch_size :].copy() if rem else None
+
+    return factory
+
+
+class DataLoader:
+    """Turn an out-of-core dataset into fixed-size minibatches for variational inference.
+
+    Parameters
+    ----------
+    dataset : iterable of ndarray, or zero-arg factory
+        The source of rows. A factory is preferred: it restarts the stream each
+        epoch. It may yield single samples or blocks of any size.
+    batch_size : int
+        Leading dimension of every yielded minibatch.
+    shuffle : bool, default False
+        Wrap the source in a bounded :func:`shuffle_buffer`.
+    buffer_size : int, optional
+        Shuffle-buffer size in rows when ``shuffle=True``; defaults to
+        ``50 * batch_size``.
+    seed : int, optional
+        Seed for the shuffle buffer (ignored when ``shuffle=False``).
+    total_size : int or "auto", default "auto"
+        The dataset size ``N``, or ``"auto"`` to infer it (from the source's
+        ``n_rows`` if available, else one counting pass). ``None`` disables
+        the rescaling.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        loader = DataLoader(
+            parquet_source("shuffled/"),
+            batch_size=4096,
+            total_size="auto",
+        )
+
+        with pm.Model() as model:
+            b = pm.Normal("b", 0.0, 3.0, shape=4)
+            batch = pm.Data("batch", np.zeros((4096, 4)))
+            logit = b[0] + b[1] * batch[:, 0] + b[2] * batch[:, 1] + b[3] * batch[:, 2]
+            pm.Bernoulli("y", logit_p=logit, observed=batch[:, 3], total_size=loader.total_size)
+
+        with model:
+            for next_batch in loader:
+                model.set_data("batch", next_batch)
+                ...
+    """
+
+    def __init__(
+        self,
+        dataset: Iterable[np.ndarray] | Callable[[], Iterator[np.ndarray]],
+        *,
+        batch_size: int,
+        shuffle: bool = False,
+        buffer_size: int | None = None,
+        seed: int | None = None,
+        total_size: int | str | None = "auto",
+    ):
+        self._new_iter = new_iter = _as_source(dataset)
+        self._batch_size = int(batch_size)
+
+        if total_size == "auto":
+            total_size = _auto_total_size(dataset, new_iter)
+        elif total_size is None:
+            warnings.warn(
+                "DataLoader created with total_size=None: the minibatch "
+                "log-likelihood will not be rescaled and the posterior will be "
+                "biased. Pass total_size=N (the true dataset size) or total_size='auto'.",
+                UserWarning,
+                stacklevel=2,
+            )
+        self._total_size = None if total_size is None else int(total_size)
+
+        if shuffle:
+            if buffer_size is None:
+                buffer_size = 50 * self._batch_size
+            self._batch_source = shuffle_buffer(
+                self._new_iter, buffer_size=buffer_size, batch_size=self._batch_size, seed=seed
+            )
+        else:
+            self._batch_source = self._new_iter
+
+    @property
+    def batch_size(self) -> int:
+        return self._batch_size
+
+    @property
+    def total_size(self) -> int | None:
+        """The dataset size ``N`` (pass to the distribution's ``total_size``)."""
+        return self._total_size
+
+    def __iter__(self) -> Iterator[np.ndarray]:
+        """Yield one epoch of ``batch_size``-row minibatches."""
+        yield from self._batch_source()
+
+    def __len__(self) -> int:
+        """Number of batches per epoch."""
+        if self._total_size is None:
+            raise TypeError(
+                "len(DataLoader) requires total_size; "
+                "construct with total_size=N or total_size='auto'."
+            )
+        return self._total_size // self._batch_size
+
+
+def _check_columns(schema, columns: list[str], path: str) -> None:
+    """Reject a shard whose schema cannot supply ``columns`` as a float batch."""
+    import pyarrow as pa
+
+    missing = [c for c in columns if c not in schema.names]
+    if missing:
+        raise ValueError(
+            f"columns {missing} not found in {path!r}; available: {sorted(schema.names)}"
+        )
+    numeric = (pa.types.is_integer, pa.types.is_floating, pa.types.is_boolean)
+    bad = [c for c in columns if not any(t(schema.field(c).type) for t in numeric)]
+    if bad:
+        raise ValueError(
+            f"columns {bad} in {path!r} are not numeric and cannot be streamed into a "
+            f"float batch; select numeric columns with columns=."
+        )
+
+
+class _ParquetDataset:
+    def __init__(self, paths: list[str], columns: list[str], n_rows: int):
+        self._paths = paths
+        self._columns = columns
+        self.n_rows = n_rows
+
+    def __iter__(self) -> Iterator[np.ndarray]:
+        import pyarrow.parquet as pq
+
+        for path in self._paths:
+            file = pq.ParquetFile(path)
+            # parquet_source only ever sees shard 0, so re-check every shard here.
+            _check_columns(file.schema_arrow, self._columns, path)
+            for i in range(file.metadata.num_row_groups):
+                table = file.read_row_group(i, columns=self._columns)
+                # Stack by the frozen names: a permuted shard must not swap features.
+                yield np.column_stack([table.column(c).to_numpy() for c in self._columns])
+
+
+def parquet_source(
+    directory: str,
+    *,
+    columns: list[str] | None = None,
+    pattern: str = "*.parquet",
+) -> _ParquetDataset:
+    """A re-iterable source over a directory of Parquet files.
+
+    Yields one ``(rows, n_columns)`` array per row group. Carries ``n_rows`` from
+    Parquet metadata so ``total_size="auto"`` is free.
+    """
+    import pyarrow.parquet as pq
+
+    paths = sorted(glob.glob(os.path.join(directory, pattern)))
+    if not paths:
+        raise ValueError(f"no Parquet files match {os.path.join(directory, pattern)!r}")
+    schema = pq.read_schema(paths[0])
+    if columns is None:
+        columns = list(schema.names)
+    _check_columns(schema, columns, paths[0])
+    n_rows = sum(pq.read_metadata(p).num_rows for p in paths)
+    return _ParquetDataset(paths, columns, n_rows)
diff --git a/tests/variational/__init__.py b/tests/variational/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/variational/dataloader_helpers.py b/tests/variational/dataloader_helpers.py
new file mode 100644
index 000000000..8c4d31e33
--- /dev/null
+++ b/tests/variational/dataloader_helpers.py
@@ -0,0 +1,70 @@
+#   Copyright 2026 - 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 DataLoader tests."""
+
+import numpy as np
+import pytest
+
+
+def chunked_factory(data, size):
+    """A zero-arg factory replaying ``data`` in ``size``-row chunks, fresh each epoch."""
+
+    def factory():
+        for i in range(0, len(data), size):
+            yield data[i : i + size]
+
+    return factory
+
+
+class BlockDataset:
+    """A re-iterable source replaying ``data`` in ``size``-row blocks.
+
+    ``n_rows`` is ``None`` unless one is given, so a dataset that knows its size
+    and one that does not are both reachable.
+    """
+
+    n_rows = None
+
+    def __init__(self, data, size, n_rows=None):
+        self._data = data
+        self._size = size
+        if n_rows is not None:
+            self.n_rows = n_rows
+
+    def __iter__(self):
+        for i in range(0, len(self._data), self._size):
+            yield self._data[i : i + self._size]
+
+
+def reused_buffer_factory(n_blocks, rows):
+    """A factory that refills and re-yields one array, as an out-of-core reader may.
+
+    Block ``i`` is ``rows`` copies of ``i``, so a batch that aliases an overwritten
+    buffer shows up as a repeated or missing value.
+    """
+
+    def factory():
+        buf = np.empty((rows, 1))
+        for value in range(n_blocks):
+            buf.fill(value)
+            yield buf
+
+    return factory
+
+
+def write_parquet(path, columns, **kwargs):
+    """Write ``columns`` to ``path``, skipping the calling test if pyarrow is not installed."""
+    pa = pytest.importorskip("pyarrow")
+    pq = pytest.importorskip("pyarrow.parquet")
+    pq.write_table(pa.table(columns), str(path), **kwargs)
diff --git a/tests/variational/test_dataloader.py b/tests/variational/test_dataloader.py
new file mode 100644
index 000000000..e4356b1c2
--- /dev/null
+++ b/tests/variational/test_dataloader.py
@@ -0,0 +1,314 @@
+#   Copyright 2026 - 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.
+import numpy as np
+import pymc as pm
+import pytest
+
+from pymc_extras.variational.dataloader import (
+    DataLoader,
+    shuffle_buffer,
+)
+from tests.variational.dataloader_helpers import (
+    BlockDataset,
+    chunked_factory,
+    reused_buffer_factory,
+)
+
+
+def test_shuffle_buffer_copies_blocks_it_holds():
+    """shuffle_buffer fills across several pulls, so it cannot alias the source's buffer."""
+    src = shuffle_buffer(reused_buffer_factory(4, 2), buffer_size=8, batch_size=4, seed=0)
+    np.testing.assert_array_equal(
+        np.sort(np.concatenate(list(src())).ravel()), np.repeat(np.arange(4, dtype="float64"), 2)
+    )
+
+
+@pytest.mark.parametrize(
+    "buffer_size, n_rows",
+    [(55, 140), (3, 120)],
+    ids=["non-dividing-chunks", "buffer-below-batch"],
+)
+def test_shuffle_buffer_conserves_rows(buffer_size, n_rows):
+    """Buffer sizes that do not divide batch_size, one of them below it, lose no rows."""
+    data = np.arange(n_rows, dtype="float64").reshape(n_rows, 1)
+    src = shuffle_buffer(chunked_factory(data, 7), buffer_size=buffer_size, batch_size=10, seed=0)
+    batches = list(src())
+    assert batches
+    assert all(b.shape == (10, 1) for b in batches)
+    seen = np.sort(np.concatenate([b.ravel() for b in batches]))
+    np.testing.assert_array_equal(seen, data.ravel())
+
+
+@pytest.mark.parametrize("chunk", [25, 100], ids=["several-chunks-per-fill", "one-chunk-fills-it"])
+def test_shuffle_buffer_does_not_mutate_source(chunk):
+    """Shuffling happens on an owned copy, even when a single chunk fills the buffer."""
+    data = np.arange(100, dtype="float64").reshape(100, 1)
+    original = data.copy()
+    src = shuffle_buffer(chunked_factory(data, chunk), buffer_size=40, batch_size=10, seed=1)
+    list(src())
+    np.testing.assert_array_equal(data, original)
+
+
+@pytest.mark.parametrize(
+    "data, make_source, batch_size, buffer_size, shapes",
+    [
+        (
+            np.arange(120.0).reshape(120, 1),
+            lambda d: chunked_factory(d, 8),
+            10,
+            40,
+            [(10, 1)] * 12,
+        ),
+        (np.arange(40.0).reshape(20, 2), lambda d: d, 8, 16, [(8, 2)] * 2),
+        (np.arange(12.0), lambda d: d, 4, 6, [(4,)] * 3),
+    ],
+    ids=["block-factory", "raw-2d", "raw-1d-scalars"],
+)
+def test_shuffle_true_yields_whole_source_rows(data, make_source, batch_size, buffer_size, shapes):
+    """shuffle=True yields only full batches, of distinct rows that all came from the source."""
+    ds = DataLoader(
+        make_source(data),
+        batch_size=batch_size,
+        shuffle=True,
+        buffer_size=buffer_size,
+        seed=0,
+        total_size=len(data),
+    )
+    batches = list(ds)
+    assert [b.shape for b in batches] == shapes
+    seen = {tuple(np.atleast_1d(r)) for b in batches for r in b}
+    assert len(seen) == sum(s[0] for s in shapes)
+    assert seen <= {tuple(np.atleast_1d(r)) for r in data}
+
+
+def test_total_size_rescales_logp_like_minibatch():
+    """total_size=loader.total_size scales the observed logp by exactly N / batch_size."""
+    rng = np.random.default_rng(0)
+    N, bs = 1000, 20
+    data = rng.normal(size=(bs, 1))
+    loader = DataLoader(lambda: iter([data]), batch_size=bs, total_size=N)
+
+    with pm.Model() as scaled:
+        mu = pm.Normal("mu", 0, 1)
+        batch = pm.Data("batch", data)
+        pm.Normal("y", mu, 1, observed=batch[:, 0], total_size=loader.total_size)
+    with pm.Model() as plain:
+        mu = pm.Normal("mu", 0, 1)
+        pm.Normal("y", mu, 1, observed=data[:, 0])
+
+    point = {"mu": np.array(0.3)}
+    obs_scaled = scaled.compile_logp(scaled.observed_RVs)(point)
+    obs_plain = plain.compile_logp(plain.observed_RVs)(point)
+    np.testing.assert_allclose(obs_scaled, obs_plain * (N / bs), rtol=1e-6)
+
+
+def test_len_raises_when_total_size_none():
+    """total_size=None warns at construction, and len() raises."""
+    data = np.ones((4, 1))
+    with pytest.warns(UserWarning, match="total_size=None"):
+        loader = DataLoader(lambda: iter([data] * 5), batch_size=4, total_size=None)
+    with pytest.raises(TypeError, match=r"len\(DataLoader\) requires total_size"):
+        len(loader)
+
+
+def test_iter_yields_clean_batches_and_reiterates():
+    """__iter__ yields batch_size-row batches and can be re-iterated."""
+    data = np.arange(40, dtype="float64").reshape(40, 1)
+    loader = DataLoader(chunked_factory(data, 10), batch_size=10, total_size=40)
+    e1 = list(loader)
+    e2 = list(loader)
+    assert len(e1) == 4 and all(b.shape == (10, 1) for b in e1)
+    np.testing.assert_array_equal(np.sort(np.concatenate([b.ravel() for b in e1])), data.ravel())
+    np.testing.assert_array_equal(np.sort(np.concatenate([b.ravel() for b in e2])), data.ravel())
+
+
+def test_accepts_numpy_integer_sizes():
+    """numpy ints pass the Integral check and are stored as plain Python ints."""
+    data = np.zeros((8, 1))
+    ds = DataLoader(chunked_factory(data, 4), batch_size=np.int64(4), total_size=np.int64(8))
+    assert next(iter(ds)).shape == (4, 1)
+    assert (ds.batch_size, ds.total_size) == (4, 8)
+    assert type(ds.batch_size) is int and type(ds.total_size) is int
+
+
+def test_shuffle_buffer_reshuffles_each_epoch_and_is_seed_reproducible():
+    """Each epoch draws a fresh permutation, and the same seed replays the same epochs."""
+    data = np.arange(60, dtype="float64").reshape(60, 1)
+
+    def epochs():
+        f = shuffle_buffer(chunked_factory(data, 10), buffer_size=60, batch_size=10, seed=7)
+        return [np.concatenate([b.ravel() for b in f()]) for _ in range(2)]
+
+    first, second = epochs()
+    assert not np.array_equal(first, second)
+    np.testing.assert_array_equal(epochs(), [first, second])
+    np.testing.assert_array_equal(np.sort(first), data.ravel())
+    np.testing.assert_array_equal(np.sort(second), data.ravel())
+
+
+def test_factory_returning_reiterable_is_accepted():
+    """A zero-arg factory may return any iterable, e.g. a list, not just an iterator."""
+    data = [np.zeros((4, 1), dtype="float64")]
+    ds = DataLoader(lambda: data, batch_size=4, total_size=4)
+    assert next(iter(ds)).shape == (4, 1)
+
+
+@pytest.mark.parametrize(
+    "data, batch_size, shapes",
+    [
+        (np.arange(6, dtype="float64"), 3, [(3,), (3,)]),
+        (np.arange(12, dtype="float64"), 4, [(4,)] * 3),
+    ],
+    ids=["exact", "drop-last"],
+)
+def test_a_one_dimensional_source_is_a_block_of_scalars(data, batch_size, shapes):
+    """A 1-D array has no trailing axis, so each element is one scalar sample."""
+    ds = DataLoader(
+        data,
+        batch_size=batch_size,
+        shuffle=True,
+        buffer_size=batch_size * 2,
+        seed=0,
+        total_size=data.size,
+    )
+    batches = list(ds)
+    assert [b.shape for b in batches] == shapes
+    np.testing.assert_array_equal(
+        np.sort(np.concatenate(batches)), np.sort(data.ravel()[: len(shapes) * batch_size])
+    )
+
+
+def test_raw_2d_array_is_one_block_of_rows():
+    """A raw 2-D array is handed over whole; with shuffle=True it's batched correctly."""
+    data = np.arange(40, dtype="float64").reshape(20, 2)
+    with pytest.warns(UserWarning, match="counting pass"):
+        ds = DataLoader(data, batch_size=8, shuffle=True, buffer_size=16, seed=0, total_size="auto")
+    assert ds.total_size == 20
+    batches = list(ds)
+    assert [b.shape for b in batches] == [(8, 2), (8, 2)]
+    streamed = np.concatenate(batches)
+    assert len({tuple(r) for r in streamed}) == 16
+    assert {tuple(r) for r in streamed} <= {tuple(r) for r in data}
+
+
+def test_shuffle_buffer_accepts_factory_returning_reiterable():
+    """A factory returning a re-iterable is normalized to one iterator, so fills don't restart."""
+    data = np.arange(120, dtype="float64").reshape(120, 1)
+    chunks = [data[i : i + 20] for i in range(0, 120, 20)]
+    src = shuffle_buffer(lambda: chunks, buffer_size=50, batch_size=10, seed=0)
+    batches = list(src())
+    assert len(batches) == 12
+    np.testing.assert_array_equal(
+        np.sort(np.concatenate([b.ravel() for b in batches])), data.ravel()
+    )
+
+
+@pytest.mark.parametrize(
+    "shuffle, buffer_size",
+    [(True, 7), (True, 1), (True, 500)],
+    ids=["buffer-under-chunk", "buffer-of-one", "buffer-over-dataset"],
+)
+@pytest.mark.parametrize(
+    "n, batch_size, chunk",
+    [(60, 10, 7), (37, 5, 1), (100, 100, 13), (23, 4, 23), (12, 1, 5)],
+    ids=[
+        "chunks-straddle-batches",
+        "one-row-chunks",
+        "one-batch-is-everything",
+        "one-chunk-many-batches",
+        "batch-of-one",
+    ],
+)
+def test_one_epoch_conserves_source_rows(n, batch_size, chunk, shuffle, buffer_size):
+    """An epoch is floor(N/batch_size) batches of distinct source rows, in source order unshuffled.
+
+    Nothing may be duplicated or invented, and the only rows a pass is allowed to
+    drop are the N mod batch_size that cannot fill a last batch -- for every mix of
+    chunking, batch size and buffer size, not just where they divide each other.
+    """
+    data = np.arange(2 * n, dtype="float64").reshape(n, 2)
+    loader = DataLoader(
+        chunked_factory(data, chunk),
+        batch_size=batch_size,
+        shuffle=shuffle,
+        buffer_size=buffer_size,
+        seed=0,
+        total_size=n,
+    )
+    batches = list(loader)
+    n_batches = n // batch_size
+    assert [b.shape for b in batches] == [(batch_size, 2)] * n_batches
+    streamed = np.concatenate(batches)
+    rows = {tuple(r) for r in streamed}
+    assert len(rows) == n_batches * batch_size
+    assert rows <= {tuple(r) for r in data}
+
+
+def test_second_epoch_replays_the_same_rows():
+    """A second pass over a re-readable source streams the same rows, reordered each epoch."""
+    data = np.arange(120, dtype="float64").reshape(60, 2)
+    loader = DataLoader(
+        chunked_factory(data, 7),
+        batch_size=10,
+        shuffle=True,
+        buffer_size=25,
+        seed=4,
+        total_size=60,
+    )
+    first, second = (np.concatenate(list(loader)) for _ in range(2))
+    for epoch in (first, second):
+        np.testing.assert_array_equal(epoch[np.argsort(epoch[:, 0])], data)
+    assert not np.array_equal(first, second)
+
+
+@pytest.mark.parametrize(
+    "buffer_size, effective", [(25, 25), (None, 500)], ids=["given", "default"]
+)
+def test_shuffled_loader_is_a_seeded_shuffle_buffer(buffer_size, effective):
+    """A seeded loader reproduces itself across instances and equals the same shuffle_buffer wrap."""
+    data = np.arange(120, dtype="float64").reshape(60, 2)
+
+    def stream(seed):
+        loader = DataLoader(
+            chunked_factory(data, 7),
+            batch_size=10,
+            shuffle=True,
+            buffer_size=buffer_size,
+            seed=seed,
+            total_size=60,
+        )
+        return np.concatenate(list(loader))
+
+    manual = shuffle_buffer(chunked_factory(data, 7), buffer_size=effective, batch_size=10, seed=11)
+    np.testing.assert_array_equal(stream(11), stream(11))
+    np.testing.assert_array_equal(stream(11), np.concatenate(list(manual())))
+    assert not np.array_equal(stream(11), stream(12))
+
+
+@pytest.mark.parametrize(
+    "make_source",
+    [
+        lambda d: chunked_factory(d, 8),
+        lambda d: [d[i : i + 8] for i in range(0, len(d), 8)],
+        lambda d: BlockDataset(d, 8),
+    ],
+    ids=["factory", "reiterable-blocks", "block-dataset"],
+)
+def test_every_source_kind_streams_the_same_batches(make_source):
+    """A block factory, a re-iterable and a BlockDataset are interchangeable."""
+    data = np.arange(80, dtype="float64").reshape(40, 2)
+    loader = DataLoader(make_source(data), batch_size=8, total_size=40)
+    batches = list(loader)
+    assert [b.shape for b in batches] == [(8, 2)] * 5
+    np.testing.assert_array_equal(np.concatenate(batches), data[:40])
diff --git a/tests/variational/test_dataloader_autosize.py b/tests/variational/test_dataloader_autosize.py
new file mode 100644
index 000000000..a3ea72a51
--- /dev/null
+++ b/tests/variational/test_dataloader_autosize.py
@@ -0,0 +1,205 @@
+#   Copyright 2026 - 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.
+"""total_size='auto' resolution, the total_size sanity warning, and parquet_source."""
+
+from contextlib import nullcontext
+
+import numpy as np
+import pytest
+
+from pymc_extras.variational.dataloader import (
+    DataLoader,
+    parquet_source,
+)
+from tests.variational.dataloader_helpers import BlockDataset, chunked_factory, write_parquet
+
+
+def test_auto_counts_finite_source():
+    """Without .n_rows, 'auto' does one counting pass and resolves the true N."""
+    data = np.arange(60, dtype="float64").reshape(60, 1)
+    with pytest.warns(UserWarning, match="counting pass"):
+        ds = DataLoader(chunked_factory(data, 7), batch_size=10, total_size="auto")
+    assert ds.total_size == 60
+
+
+@pytest.mark.parametrize("shuffle", [False, True], ids=["plain", "shuffled"])
+def test_auto_uses_n_rows_fast_path(shuffle):
+    """A source-advertised .n_rows is trusted without a counting pass, shuffle or not."""
+    src = chunked_factory(np.zeros((8, 1)), 4)
+    src.n_rows = 1000
+    ds = DataLoader(
+        src,
+        batch_size=4,
+        shuffle=shuffle,
+        buffer_size=8,
+        seed=0,
+        total_size="auto",
+    )
+    assert ds.total_size == 1000
+
+
+def test_auto_rejects_one_shot_iterator():
+    """A bare generator would be consumed by the counting pass, so 'auto' refuses it."""
+    data = np.zeros((20, 1))
+    one_shot = (data[i : i + 4] for i in range(0, 20, 4))
+    with pytest.raises(ValueError, match="re-readable"):
+        DataLoader(one_shot, batch_size=4, total_size="auto")
+
+
+def test_auto_rejects_factory_returning_same_one_shot_iterator():
+    """A factory handing back the same consumed iterator is not re-readable."""
+    data = np.zeros((20, 1))
+    one_shot = (data[i : i + 4] for i in range(0, 20, 4))
+    with (
+        pytest.warns(UserWarning, match="counting pass"),
+        pytest.raises(ValueError, match="fresh iterator"),
+    ):
+        DataLoader(lambda: one_shot, batch_size=4, total_size="auto")
+
+
+def test_auto_accepts_any_n_rows():
+    """A source .n_rows is trusted as-is (like PyTorch)."""
+    f = chunked_factory(np.zeros((8, 1)), 4)
+    f.n_rows = 0
+    ds = DataLoader(f, batch_size=4, total_size="auto")
+    assert ds.total_size == 0
+
+
+def test_auto_rejects_factory_closing_over_consumed_iterator():
+    """A generator function over a consumed iterator returns a new but empty stream."""
+    data = np.zeros((20, 1))
+    underlying = iter([data[i : i + 4] for i in range(0, 20, 4)])
+
+    def gen():
+        yield from underlying
+
+    with (
+        pytest.warns(UserWarning, match="counting pass"),
+        pytest.raises(ValueError, match="fresh iterator"),
+    ):
+        DataLoader(gen, batch_size=4, total_size="auto")
+
+
+def test_parquet_source_n_rows_from_metadata(tmp_path):
+    """n_rows comes from file metadata, and total_size='auto' takes it without counting."""
+    rng = np.random.default_rng(0)
+    total = 0
+    for i in range(3):
+        n = 100 + 50 * i
+        total += n
+        block = rng.normal(size=(n, 2))
+        write_parquet(tmp_path / f"part_{i:02d}.parquet", {"a": block[:, 0], "b": block[:, 1]})
+    src = parquet_source(str(tmp_path))
+    assert src.n_rows == total
+
+    ds = DataLoader(src, batch_size=10, total_size="auto")
+    assert ds.total_size == total
+
+
+def test_parquet_source_columns_and_shard_order(tmp_path):
+    """columns= selects a column subset and shards are read in sorted path order."""
+    for i in range(2):
+        write_parquet(
+            tmp_path / f"part_{i}.parquet",
+            {"a": [float(i)] * 2, "b": [9.0] * 2, "c": [float(10 + i)] * 2},
+        )
+    src = parquet_source(str(tmp_path), columns=["a", "c"])
+    blocks = list(src)
+    assert [b.shape for b in blocks] == [(2, 2), (2, 2)]
+    np.testing.assert_array_equal(blocks[0][:, 0], [0.0, 0.0])
+    np.testing.assert_array_equal(blocks[1][:, 1], [11.0, 11.0])
+
+
+def test_parquet_source_empty_dir_raises(tmp_path):
+    """A directory with no matching Parquet files raises a clear error."""
+    pytest.importorskip("pyarrow")
+    with pytest.raises(ValueError, match="no Parquet files match"):
+        parquet_source(str(tmp_path))
+
+
+def test_parquet_source_freezes_column_order_across_permuted_shards(tmp_path):
+    """A shard whose schema permutes the columns is read back in the first shard's order."""
+    write_parquet(tmp_path / "p0.parquet", {"a": [1.0, 1.0], "b": [10.0, 10.0]})
+    write_parquet(tmp_path / "p1.parquet", {"b": [20.0, 20.0], "a": [2.0, 2.0]})
+    blocks = list(parquet_source(str(tmp_path)))
+    np.testing.assert_array_equal(blocks[0], [[1.0, 10.0], [1.0, 10.0]])
+    np.testing.assert_array_equal(blocks[1], [[2.0, 20.0], [2.0, 20.0]])
+
+
+def test_parquet_source_streams_row_groups_not_whole_files(tmp_path):
+    """A multi-row-group file is yielded one row group at a time, not one file at a time."""
+    write_parquet(tmp_path / "p.parquet", {"a": np.arange(30.0)}, row_group_size=10)
+    blocks = list(parquet_source(str(tmp_path)))
+    assert [b.shape for b in blocks] == [(10, 1), (10, 1), (10, 1)]
+    np.testing.assert_array_equal(np.concatenate(blocks).ravel(), np.arange(30.0))
+
+
+def test_parquet_source_names_a_later_shard_with_a_non_numeric_column(tmp_path):
+    """A later shard whose column turned non-numeric is named by path, not an opaque cast error."""
+    write_parquet(tmp_path / "p0.parquet", {"a": [1.0, 2.0]})
+    write_parquet(tmp_path / "p1.parquet", {"a": ["bad", "worse"]})
+    src = parquet_source(str(tmp_path))
+    with pytest.raises(ValueError, match=r"p1\.parquet.*not numeric"):
+        list(src)
+
+
+def test_parquet_source_rejects_non_numeric_columns(tmp_path):
+    """A string column is rejected at construction, naming the column and the columns= remedy."""
+    write_parquet(tmp_path / "p.parquet", {"x": [1.0, 2.0], "id": ["a", "b"]})
+    with pytest.raises(ValueError, match="not numeric"):
+        parquet_source(str(tmp_path))
+    src = parquet_source(str(tmp_path), columns=["x"])
+    np.testing.assert_array_equal(next(iter(src)), [[1.0], [2.0]])
+
+
+def test_parquet_source_names_the_shard_missing_a_column(tmp_path):
+    """read_row_group drops unknown names silently, so a shard missing a frozen column is named."""
+    write_parquet(tmp_path / "p0.parquet", {"a": [1.0], "b": [2.0]})
+    write_parquet(tmp_path / "p1.parquet", {"a": [3.0]})
+    src = parquet_source(str(tmp_path))
+    with pytest.raises(ValueError, match=r"p1\.parquet"):
+        list(src)
+
+
+def test_parquet_source_rejects_unknown_columns(tmp_path):
+    """A typo in columns= raises at construction, not as a pyarrow error at first iteration."""
+    write_parquet(tmp_path / "p.parquet", {"a": [1.0], "b": [2.0]})
+    with pytest.raises(ValueError, match="not found"):
+        parquet_source(str(tmp_path), columns=["a", "nope"])
+
+
+@pytest.mark.parametrize(
+    "make_source, counts",
+    [
+        (lambda d: chunked_factory(d, 6), True),
+        (lambda d: [d[i : i + 6] for i in range(0, len(d), 6)], True),
+        (lambda d: BlockDataset(d, 6), True),
+        (lambda d: BlockDataset(d, 6, n_rows=42), False),
+    ],
+    ids=["factory", "reiterable-blocks", "iterable-dataset", "advertised-n-rows"],
+)
+def test_auto_resolves_the_explicit_n_for_every_source_kind(make_source, counts):
+    """'auto' lands on the same N an explicit total_size would, counting only when it must.
+
+    A source that does not advertise n_rows is counted, a BlockDataset that
+    inherits the default ``n_rows=None`` is counted too, and either way the epoch that
+    follows the counting pass is the whole dataset.
+    """
+    data = np.arange(84, dtype="float64").reshape(42, 2)
+    counting = pytest.warns(UserWarning, match="counting pass") if counts else nullcontext()
+    with counting:
+        ds = DataLoader(make_source(data), batch_size=6, total_size="auto")
+    assert ds.total_size == 42
+    assert len(ds) == 42 // 6
+    np.testing.assert_array_equal(np.concatenate(list(ds)), data[:42])