# Source code for tensorplay.utils.data

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/utils/data.html

```
"""Data loading utilities: Dataset, DataLoader, Sampler, collate.

The public names, validation order, and error messages remain stable; the
implementation is built on ``tensorplay.Tensor``.

The implementation also has native batch-fetch hooks for TensorPlay-backed
the default collate path to return a single indexed batch instead of creating
one Python tensor view per sample.
"""

from __future__ import annotations

import bisect
import copy
import contextlib
import itertools
import math
import multiprocessing
import os
import queue
import random
import re
import sys
import threading
import traceback
import warnings
from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence, Sized
from typing import Any, Callable, Generic, Iterable, Iterator, List, Optional, TypeVar

import numpy as np

import tensorplay as tp
from tensorplay import Tensor

__all__ = [
    "BatchSampler",
    "ChainDataset",
    "ConcatDataset",
    "DataLoader",
    "Dataset",
    "DistributedSampler",
    "IterableDataset",
    "RandomSampler",
    "Sampler",
    "SequentialSampler",
    "StackDataset",
    "Subset",
    "SubsetRandomSampler",
    "TensorDataset",
    "WeightedRandomSampler",
    "_DatasetKind",
    "default_collate",
    "default_convert",
    "get_worker_info",
    "random_split",
]

_T_co = TypeVar("_T_co", covariant=True)
_T = TypeVar("_T")

# A private sentinel used by dataset wrappers to say that their zero-copy
# batch path cannot represent the requested batch.  The fetcher then falls back
_FAST_BATCH_UNAVAILABLE = object()

def _make_batch_index(indices, tensor: Tensor) -> Tensor:
    """Build an index tensor on the same device as ``tensor``.

    Keep batch indices in an explicit int64 tensor so the dataset fast paths
    use one ``index_select`` on both CPU and CUDA.  This avoids constructing a
    Python tensor view for every sample even though Tensor now also supports
    the public ``tensor[[...]]`` spelling.
    """
    if isinstance(indices, Tensor):
        values = indices.tolist()
    else:
        values = [int(index) for index in indices]
    return tp.tensor(values, dtype=tp.int64, device=tensor.device)

# ---------------------------------------------------------------------------
# Datasets
# ---------------------------------------------------------------------------

[docs]
class Dataset(Generic[_T_co]):
    """An abstract class representing a :class:`Dataset`.

    All datasets that represent a map from keys to data samples should subclass
    it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching
    a data sample for a given key. Subclasses could also optionally overwrite
    :meth:`__len__`, which is expected to return the size of the dataset by many
    :class:`~tensorplay.utils.data.Sampler` implementations and the default
    options of :class:`~tensorplay.utils.data.DataLoader`.

    .. note::
      :class:`~tensorplay.utils.data.DataLoader` by default constructs an index
      sampler that yields integral indices. To make it work with a map-style
      dataset with non-integral indices/keys, a custom sampler must be provided.
    """

    def __getitem__(self, index) -> _T_co:
        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")

    def __add__(self, other: "Dataset[_T_co]") -> "ConcatDataset[_T_co]":
        return ConcatDataset([self, other])

[docs]
class IterableDataset(Dataset[_T_co], Iterable[_T_co]):
    """An iterable Dataset.

    All datasets that represent an iterable of data samples should subclass it.
    Such form of datasets is particularly useful when data come from a stream.

    All subclasses should overwrite :meth:`__iter__`, which would return an
    iterator of samples in this dataset.

    When a subclass is used with :class:`~tensorplay.utils.data.DataLoader`,
    each item in the dataset will be yielded from the DataLoader iterator.
    """

    def __add__(self, other: Dataset[_T_co]):
        return ChainDataset([self, other])

[docs]
class TensorDataset(Dataset[tuple[Tensor, ...]]):
    """Dataset wrapping tensors.

    Each sample will be retrieved by indexing tensors along the first dimension.

    Args:
        *tensors (Tensor): tensors that have the same size of the first dimension.
    """

    tensors: tuple[Tensor, ...]

    def __init__(self, *tensors: Tensor) -> None:
        if any(tensors[0].size(0) != tensor.size(0) for tensor in tensors):
            raise AssertionError("Size mismatch between tensors")
        self.tensors = tensors

    def __getitem__(self, index):
        if isinstance(index, list):
            return tuple(
                tp.index_select(tensor, 0, _make_batch_index(index, tensor))
                for tensor in self.tensors
            )
        return tuple(tensor[index] for tensor in self.tensors)

    def _tp_get_batch(self, indices):
        """Return the default-collated batch without per-row Python views."""
        return [
            tp.index_select(tensor, 0, _make_batch_index(indices, tensor))
            for tensor in self.tensors
        ]

    def __getitems__(self, indices: list[int]):
        # default-collate path uses ``_tp_get_batch`` above to avoid rebuilding
        # this list and stacking it a second time.
        batch = self._tp_get_batch(indices)
        return [tuple(tensor[i] for tensor in batch) for i in range(len(indices))]

    def __len__(self) -> int:
        return self.tensors[0].size(0)

[docs]
class StackDataset(Dataset[_T]):
    """Dataset as a stacking of multiple datasets.

    This class is useful to assemble different parts of complex input data,
    given as datasets.

    Example:
        >>> images = ImageDataset()
        >>> texts = TextDataset()
        >>> tuple_stack = StackDataset(images, texts)
        >>> tuple_stack[0] == (images[0], texts[0])
        >>> dict_stack = StackDataset(image=images, text=texts)
        >>> dict_stack[0] == {"image": images[0], "text": texts[0]}

    Args:
        *args (Dataset): Datasets for stacking returned as tuple.
        **kwargs (Dataset): Datasets for stacking returned as dict.
    """

    datasets: tuple | dict

    def __init__(self, *args: Dataset[_T_co], **kwargs: Dataset[_T_co]) -> None:
        if args:
            if kwargs:
                raise ValueError(
                    "Supported either ``tuple``- (via ``args``) or"
                    "``dict``- (via ``kwargs``) like input/output, but both types are given."
                )
            self._length = len(args[0])
            if any(self._length != len(dataset) for dataset in args):
                raise ValueError("Size mismatch between datasets")
            self.datasets = args
        elif kwargs:
            tmp = list(kwargs.values())
            self._length = len(tmp[0])
            if any(self._length != len(dataset) for dataset in tmp):
                raise ValueError("Size mismatch between datasets")
            self.datasets = kwargs
        else:
            raise ValueError("At least one dataset should be passed")

    def __getitem__(self, index):
        if isinstance(self.datasets, dict):
            return {k: dataset[index] for k, dataset in self.datasets.items()}
        return tuple(dataset[index] for dataset in self.datasets)

    def __getitems__(self, indices: list):
        # add batched sampling support when parent datasets supports it.
        if isinstance(self.datasets, dict):
            dict_batch: list[dict] = [{} for _ in indices]
            for k, dataset in self.datasets.items():
                if callable(getattr(dataset, "__getitems__", None)):
                    items = dataset.__getitems__(indices)
                    if len(items) != len(indices):
                        raise ValueError(
                            "Nested dataset's output size mismatch."
                            f" Expected {len(indices)}, got {len(items)}"
                        )
                    for data, d_sample in zip(items, dict_batch, strict=True):
                        d_sample[k] = data
                else:
                    for idx, d_sample in zip(indices, dict_batch, strict=True):
                        d_sample[k] = dataset[idx]
            return dict_batch

        # tuple data
        list_batch: list[list] = [[] for _ in indices]
        for dataset in self.datasets:
            if callable(getattr(dataset, "__getitems__", None)):
                items = dataset.__getitems__(indices)
                if len(items) != len(indices):
                    raise ValueError(
                        "Nested dataset's output size mismatch."
                        f" Expected {len(indices)}, got {len(items)}"
                    )
                for data, t_sample in zip(items, list_batch, strict=True):
                    t_sample.append(data)
            else:
                for idx, t_sample in zip(indices, list_batch, strict=True):
                    t_sample.append(dataset[idx])
        tuple_batch: list[tuple] = [tuple(sample) for sample in list_batch]
        return tuple_batch

    def _tp_get_batch(self, indices):
        """Compose native batches from children when every child supports it."""
        if isinstance(self.datasets, dict):
            result = {}
            for key, dataset in self.datasets.items():
                getter = getattr(dataset, "_tp_get_batch", None)
                if not callable(getter):
                    return _FAST_BATCH_UNAVAILABLE
                batch = getter(indices)
                if batch is _FAST_BATCH_UNAVAILABLE:
                    return _FAST_BATCH_UNAVAILABLE
                result[key] = batch
            return result

        result = []
        for dataset in self.datasets:
            getter = getattr(dataset, "_tp_get_batch", None)
            if not callable(getter):
                return _FAST_BATCH_UNAVAILABLE
            batch = getter(indices)
            if batch is _FAST_BATCH_UNAVAILABLE:
                return _FAST_BATCH_UNAVAILABLE
            result.append(batch)
        return result

    def __len__(self) -> int:
        return self._length

[docs]
class ConcatDataset(Dataset[_T_co]):
    """Dataset as a concatenation of multiple datasets.

    This class is useful to assemble different existing datasets.

    Args:
        datasets (sequence): List of datasets to be concatenated
    """

    datasets: list[Dataset[_T_co]]
    cumulative_sizes: list[int]

    @staticmethod
    def cumsum(sequence):
        r, s = [], 0
        for e in sequence:
            l = len(e)
            r.append(l + s)
            s += l
        return r

    def __init__(self, datasets: Iterable[Dataset]) -> None:
        super().__init__()
        self.datasets = list(datasets)
        if len(self.datasets) == 0:
            raise AssertionError("datasets should not be an empty iterable")
        for d in self.datasets:
            if isinstance(d, IterableDataset):
                raise AssertionError("ConcatDataset does not support IterableDataset")
        self.cumulative_sizes = self.cumsum(self.datasets)

    def __len__(self) -> int:
        return self.cumulative_sizes[-1]

    def __getitem__(self, idx):
        dataset_idx, sample_idx = self._resolve_index(idx)
        return self.datasets[dataset_idx][sample_idx]

    def _resolve_index(self, idx):
        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    "absolute value of index should not exceed dataset length"
                )
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx >= len(self.datasets):
            # Match the natural IndexError raised by the scalar path while
            # giving the batch path a useful, deterministic failure.
            raise IndexError("list index out of range")
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return dataset_idx, sample_idx

    def __getitems__(self, indices: list[int]):
        """Fetch a batch while grouping requests by child dataset.

        This preserves the order of the caller's indices, but lets child
        datasets use their own batch-aware ``__getitems__`` implementation.
        It removes a Python-level binary search and dispatch for every child
        sample in the common concatenated-dataset case.
        """
        if not indices:
            return []

        grouped: dict[int, list[tuple[int, int]]] = {}
        for output_position, index in enumerate(indices):
            dataset_idx, sample_idx = self._resolve_index(index)
            grouped.setdefault(dataset_idx, []).append((output_position, sample_idx))

        result = [None] * len(indices)
        for dataset_idx, positions in grouped.items():
            child = self.datasets[dataset_idx]
            child_indices = [sample_idx for _, sample_idx in positions]
            getter = getattr(child, "__getitems__", None)
            if callable(getter):
                child_items = getter(child_indices)
                if len(child_items) != len(child_indices):
                    raise ValueError(
                        "Nested dataset's output size mismatch."
                        f" Expected {len(child_indices)}, got {len(child_items)}"
                    )
            else:
                child_items = [child[index] for index in child_indices]
            for (output_position, _), item in zip(positions, child_items, strict=True):
                result[output_position] = item
        return result

    def _tp_get_batch(self, indices):
        # A batch crossing a child boundary needs structure-aware merging.  The
        # public __getitems__ path handles that case; keep the direct native
        # path for the frequent single-child batch where it is unambiguous.
        if not indices:
            return _FAST_BATCH_UNAVAILABLE
        dataset_idx, sample_idx = self._resolve_index(indices[0])
        local_indices = [sample_idx]
        for index in indices[1:]:
            next_dataset_idx, next_sample_idx = self._resolve_index(index)
            if next_dataset_idx != dataset_idx:
                return _FAST_BATCH_UNAVAILABLE
            local_indices.append(next_sample_idx)
        getter = getattr(self.datasets[dataset_idx], "_tp_get_batch", None)
        if not callable(getter):
            return _FAST_BATCH_UNAVAILABLE
        return getter(local_indices)

    @property
    def cummulative_sizes(self):
        warnings.warn(
            "`cummulative_sizes` attribute is renamed to `cumulative_sizes`",
            FutureWarning,
            stacklevel=2,
        )
        return self.cumulative_sizes

[docs]
class ChainDataset(IterableDataset):
    """Dataset for chaining multiple :class:`IterableDataset` s.

    This class is useful to assemble different existing dataset streams. The
    chaining operation is done on-the-fly, so concatenating large-scale
    datasets with this class will be efficient.

    Args:
        datasets (iterable of IterableDataset): datasets to be chained together
    """

    def __init__(self, datasets: Iterable[Dataset]) -> None:
        super().__init__()
        self.datasets = datasets

    def __iter__(self):
        for d in self.datasets:
            if not isinstance(d, IterableDataset):
                raise AssertionError("ChainDataset only supports IterableDataset")
            yield from d

    def __len__(self) -> int:
        total = 0
        for d in self.datasets:
            if not isinstance(d, IterableDataset):
                raise AssertionError("ChainDataset only supports IterableDataset")
            total += len(d)
        return total

[docs]
class Subset(Dataset[_T_co]):
    """Subset of a dataset at specified indices.

    .. note::
        When subclassing `Subset` and overriding `__getitem__`, you **must** also
        override `__getitems__` to ensure `DataLoader` works correctly with your
        custom logic. If you override only `__getitem__`, a `NotImplementedError`
        will be raised when using `DataLoader`.

    Args:
        dataset (Dataset): The whole Dataset
        indices (sequence): Indices in the whole set selected for subset
    """

    dataset: Dataset[_T_co]
    indices: Sequence[int]

    def __init__(self, dataset: Dataset[_T_co], indices: Sequence[int]) -> None:
        self.dataset = dataset
        self.indices = indices

        # Check if __getitem__ is overridden but __getitems__ is not
        if (
            type(self).__getitem__ is not Subset.__getitem__
            and type(self).__getitems__ is Subset.__getitems__
        ):
            raise NotImplementedError(
                f"{type(self).__name__} overrides __getitem__ but not __getitems__. "
                "When subclassing Subset and overriding __getitem__, you must also override "
                "__getitems__ to ensure DataLoader works correctly with your custom logic. "
                "A simple implementation:\n\n"
                "def __getitems__(self, indices):\n"
                "    return [self.__getitem__(idx) for idx in indices]"
            )

    def __getitem__(self, idx):
        if isinstance(idx, list):
            return self.dataset[[self.indices[i] for i in idx]]
        return self.dataset[self.indices[idx]]

    def __getitems__(self, indices: list[int]) -> list[_T_co]:
        # add batched sampling support when parent dataset supports it.
        if callable(getattr(self.dataset, "__getitems__", None)):
            return self.dataset.__getitems__([self.indices[idx] for idx in indices])
        else:
            return [self.dataset[self.indices[idx]] for idx in indices]

    def _tp_get_batch(self, indices):
        # Never bypass a subclass' explicitly customized fetch contract.
        if type(self) is not Subset:
            return _FAST_BATCH_UNAVAILABLE
        getter = getattr(self.dataset, "_tp_get_batch", None)
        if not callable(getter):
            return _FAST_BATCH_UNAVAILABLE
        return getter([self.indices[idx] for idx in indices])

    def __len__(self) -> int:
        return len(self.indices)

_GENERATOR_LOCK = threading.RLock()

@contextlib.contextmanager
def _use_generator(generator: Optional["tp.Generator"]):
    """Run TensorPlay's global-generator kernels against ``generator``.

    The native random operators currently consume the process-global CPU
    generator.  Swapping its complete serialized state (and restoring it in a
    explicit-generator APIs: the supplied generator advances, while the
    process-global generator is left untouched.  A foreign generator object
    is retained as a compatibility fallback and is seeded by ``initial_seed``
    when its state tensor cannot be exchanged with TensorPlay.
    """
    if generator is None:
        yield
        return

    with _GENERATOR_LOCK:
        global_state = tp.get_rng_state()
        try:
            generator_state = generator.get_state()
            if not isinstance(generator_state, Tensor):
                raise TypeError("generator state is not a TensorPlay Tensor")
            tp.set_rng_state(generator_state)
        except (AttributeError, TypeError, RuntimeError):
            # Keep the old best-effort behavior for foreign Generator-like
            # objects.  Most callers use tp.Generator, for which the branch
            # above is fully stateful.
            tp.manual_seed(generator.initial_seed())
            try:
                yield
            finally:
                tp.set_rng_state(global_state)
            return

        try:
            yield
        finally:
            generator.set_state(tp.get_rng_state())
            tp.set_rng_state(global_state)

def random_split(
    dataset: Dataset[_T],
    lengths: Sequence[int | float],
    generator: Optional["tp.Generator"] = None,
) -> list[Subset[_T]]:
    """Randomly split a dataset into non-overlapping new datasets of given lengths.

    If a list of fractions that sum up to 1 is given, the lengths will be
    computed automatically as floor(frac * len(dataset)) for each fraction
    provided. After computing the lengths, if there are any remainders, 1 count
    will be distributed in round-robin fashion to the lengths until there are
    no remainders left.

    Args:
        dataset (Dataset): Dataset to be split
        lengths (sequence): lengths or fractions of splits to be produced
        generator (Generator): Generator used for the random permutation.
    """
    if math.isclose(sum(lengths), 1) and sum(lengths) <= 1:
        subset_lengths: list[int] = []
        for i, frac in enumerate(lengths):
            if frac < 0 or frac > 1:
                raise ValueError(f"Fraction at index {i} is not between 0 and 1")
            n_items_in_split = math.floor(len(dataset) * frac)
            subset_lengths.append(n_items_in_split)
        remainder = len(dataset) - sum(subset_lengths)
        # add 1 to all the lengths in round-robin fashion until the remainder is 0
        for i in range(remainder):
            idx_to_add_at = i % len(subset_lengths)
            subset_lengths[idx_to_add_at] += 1
        lengths = subset_lengths
        for i, length in enumerate(lengths):
            if length == 0:
                warnings.warn(
                    f"Length of split at index {i} is 0. "
                    f"This might result in an empty dataset.",
                    stacklevel=2,
                )

    # Cannot verify that dataset is Sized
    if sum(lengths) != len(dataset):
        raise ValueError(
            "Sum of input lengths does not equal the length of the input dataset!"
        )

    with _use_generator(generator):
        indices = tp.randperm(sum(lengths)).tolist()
    return [
        Subset(dataset, indices[offset - length : offset])
        for offset, length in zip(itertools.accumulate(lengths), lengths, strict=True)
    ]

# ---------------------------------------------------------------------------
# Samplers
# ---------------------------------------------------------------------------

[docs]
class Sampler(Generic[_T_co]):
    r"""Base class for all Samplers.

    Every Sampler subclass has to provide an :meth:`__iter__` method, providing
    a way to iterate over indices or lists of indices (batches) of dataset
    elements, and may provide a :meth:`__len__` method that returns the length
    of the returned iterators.

    .. note:: The :meth:`__len__` method isn't strictly required by
              :class:`~tensorplay.utils.data.DataLoader`, but is expected in any
              calculation involving the length of a DataLoader.
    """

    def __iter__(self) -> Iterator[_T_co]:
        raise NotImplementedError

[docs]
class SequentialSampler(Sampler[int]):
    """Samples elements sequentially, always in the same order.

    Args:
        data_source (Sized): data source to sample from. Must implement __len__.
    """

    data_source: Sized

    def __init__(self, data_source: Sized) -> None:
        self.data_source = data_source

    def __iter__(self) -> Iterator[int]:
        return iter(range(len(self.data_source)))

    def __len__(self) -> int:
        return len(self.data_source)

[docs]
class RandomSampler(Sampler[int]):
    """Samples elements randomly. If without replacement, then sample from a
    shuffled dataset. If with replacement, then user can specify
    :attr:`num_samples` to draw.

    Args:
        data_source (Sized): data source to sample from. Must implement __len__.
        replacement (bool): samples are drawn on-demand with replacement if
            ``True``, default=``False``.
        num_samples (int): number of samples to draw, default=`len(dataset)`.
        generator (Generator): Generator used in sampling.
    """

    data_source: Sized
    replacement: bool

    def __init__(
        self,
        data_source: Sized,
        replacement: bool = False,
        num_samples: Optional[int] = None,
        generator: Optional["tp.Generator"] = None,
    ) -> None:
        self.data_source = data_source
        self.replacement = replacement
        self._num_samples = num_samples
        self.generator = generator

        if not isinstance(self.replacement, bool):
            raise TypeError(
                f"replacement should be a boolean value, but got replacement={self.replacement}"
            )

        if not isinstance(self.num_samples, int) or self.num_samples <= 0:
            raise ValueError(
                f"num_samples should be a positive integer value, but got num_samples={self.num_samples}"
            )

    @property
    def num_samples(self) -> int:
        # dataset size might change at runtime
        if self._num_samples is None:
            return len(self.data_source)
        return self._num_samples

    def __iter__(self) -> Iterator[int]:
        n = len(self.data_source)
        if self.generator is None:
            # Fresh random seed per epoch, drawn from the global generator
            seed = int(tp.empty((), dtype=tp.int64).random_(0, 2**63 - 1).item())
            generator = tp.Generator()
            generator.manual_seed(seed)
        else:
            generator = self.generator

        # Materialize before yielding so the global-generator swap never
        # remains active while a caller pauses a partially consumed sampler.
        with _use_generator(generator):
            if self.replacement:
                indices = []
                for _ in range(self.num_samples // 32):
                    indices.extend(tp.randint(0, n, (32,)).tolist())
                indices.extend(tp.randint(0, n, (self.num_samples % 32,)).tolist())
            else:
                indices = []
                for _ in range(self.num_samples // n):
                    indices.extend(tp.randperm(n).tolist())
                indices.extend(tp.randperm(n).tolist()[: self.num_samples % n])
        yield from indices

    def __len__(self) -> int:
        return self.num_samples

[docs]
class SubsetRandomSampler(Sampler[int]):
    """Samples elements randomly from a given list of indices, without replacement.

    Args:
        indices (sequence): a sequence of indices
        generator (Generator): Generator used in sampling.
    """

    indices: Sequence[int]

    def __init__(self, indices: Sequence[int], generator: Optional["tp.Generator"] = None) -> None:
        self.indices = indices
        self.generator = generator

    def __iter__(self) -> Iterator[int]:
        with _use_generator(self.generator):
            permutation = tp.randperm(len(self.indices)).tolist()
        yield from (self.indices[i] for i in permutation)

    def __len__(self) -> int:
        return len(self.indices)

[docs]
class WeightedRandomSampler(Sampler[int]):
    """Samples elements from ``[0,..,len(weights)-1]`` with given probabilities (weights).

    Args:
        weights (sequence): a sequence of weights, not necessary summing up to one
        num_samples (int): number of samples to draw
        replacement (bool): if ``True``, samples are drawn with replacement.
            If not, they are drawn without replacement, which means that when a
            sample index is drawn for a row, it cannot be drawn again for that row.
        generator (Generator): Generator used in sampling.
    """

    weights: Tensor
    num_samples: int
    replacement: bool

    def __init__(
        self,
        weights: Sequence[float],
        num_samples: int,
        replacement: bool = True,
        generator: Optional["tp.Generator"] = None,
    ) -> None:
        if (
            not isinstance(num_samples, int)
            or isinstance(num_samples, bool)
            or num_samples <= 0
        ):
            raise ValueError(
                f"num_samples should be a positive integer value, but got num_samples={num_samples}"
            )
        if not isinstance(replacement, bool):
            raise ValueError(
                f"replacement should be a boolean value, but got replacement={replacement}"
            )

        weights_tensor = tp.as_tensor(np.asarray(weights, dtype=np.float64))
        if len(weights_tensor.shape) != 1:
            raise ValueError(
                "weights should be a 1d sequence but given "
                f"weights have shape {tuple(weights_tensor.shape)}"
            )

        self.weights = weights_tensor
        self.num_samples = num_samples
        self.replacement = replacement
        self.generator = generator

    def __iter__(self) -> Iterator[int]:
        with _use_generator(self.generator):
            rand_tensor = tp.multinomial(
                self.weights, self.num_samples, self.replacement
            )
            indices = rand_tensor.tolist()
        yield from indices

    def __len__(self) -> int:
        return self.num_samples

[docs]
class BatchSampler(Sampler[list[int]]):
    """Wraps another sampler to yield a mini-batch of indices.

    Args:
        sampler (Sampler or Iterable): Base sampler. Can be any iterable object
        batch_size (int): Size of mini-batch.
        drop_last (bool): If ``True``, the sampler will drop the last batch if
            its size would be less than ``batch_size``
    """

    def __init__(
        self,
        sampler: Sampler[int] | Iterable[int],
        batch_size: int,
        drop_last: bool,
    ) -> None:
        # Since collections.abc.Iterable does not check for `__getitem__`, which
        # is one way for an object to be an iterable, we don't do an `isinstance`
        # check here.
        if (
            not isinstance(batch_size, int)
            or isinstance(batch_size, bool)
            or batch_size <= 0
        ):
            raise ValueError(
                f"batch_size should be a positive integer value, but got batch_size={batch_size}"
            )
        if not isinstance(drop_last, bool):
            raise ValueError(
                f"drop_last should be a boolean value, but got drop_last={drop_last}"
            )
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last

    def __iter__(self) -> Iterator[list[int]]:
        sampler_iter = iter(self.sampler)
        if self.drop_last:
            # Create multiple references to the same iterator
            args = [sampler_iter] * self.batch_size
            for batch_droplast in zip(*args, strict=False):
                yield [*batch_droplast]
        else:
            batch = [*itertools.islice(sampler_iter, self.batch_size)]
            while batch:
                yield batch
                batch = [*itertools.islice(sampler_iter, self.batch_size)]

    def __len__(self) -> int:
        # Can only be called if self.sampler has __len__ implemented
        if self.drop_last:
            return len(self.sampler) // self.batch_size
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size

# ---------------------------------------------------------------------------
# Distributed sampling
# ---------------------------------------------------------------------------
def _dist_world_size_and_rank():
    """Best-effort world size / rank from tensorplay.distributed."""
    try:
        import tensorplay.distributed as dist

        if dist.is_available() and dist.is_initialized():
            return dist.get_world_size(), dist.get_rank()
    except ImportError:
        pass
    raise RuntimeError("Requires distributed package to be available")

class DistributedSampler(Sampler[_T_co]):
    r"""Sampler that restricts data loading to a subset of the dataset.

    It is especially useful in conjunction with distributed data parallel
    training: each process passes a :class:`DistributedSampler` as its
    :class:`~tensorplay.utils.data.DataLoader` sampler and loads an exclusive
    subset of the dataset.

    .. note::
        Dataset is assumed to be of constant size and that any instance of it
        always returns the same elements in the same order.

    Args:
        dataset: Dataset used for sampling.
        num_replicas (int, optional): Number of processes participating in
            distributed training. By default, :attr:`world_size` is retrieved
            from the current distributed group.
        rank (int, optional): Rank of the current process within
            :attr:`num_replicas`. By default, :attr:`rank` is retrieved from
            the current distributed group.
        shuffle (bool, optional): If ``True`` (default), sampler will shuffle
            the indices.
        seed (int, optional): random seed used to shuffle the sampler if
            :attr:`shuffle=True`. This number should be identical across all
            processes in the distributed group. Default: ``0``.
        drop_last (bool, optional): if ``True``, then the sampler will drop
            the tail of the data to make it evenly divisible across the number
            of replicas. If ``False``, the sampler will add extra indices to
            make the data evenly divisible across the replicas. Default:
            ``False``.

    .. warning::
        In distributed mode, calling the :meth:`set_epoch` method at the
        beginning of each epoch **before** creating the DataLoader iterator is
        necessary to make shuffling work properly across multiple epochs.
        Otherwise, the same ordering will be always used.
    """

    def __init__(
        self,
        dataset: Dataset[_T_co],
        num_replicas: Optional[int] = None,
        rank: Optional[int] = None,
        shuffle: bool = True,
        seed: int = 0,
        drop_last: bool = False,
    ) -> None:
        if num_replicas is None or rank is None:
            default_world_size, default_rank = _dist_world_size_and_rank()
            if num_replicas is None:
                num_replicas = default_world_size
            if rank is None:
                rank = default_rank
        if rank >= num_replicas or rank < 0:
            raise ValueError(
                f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]"
            )
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        # If the dataset length is evenly divisible by # of replicas, then
        # there is no need to drop any data, since the dataset will be split
        # equally.
        if self.drop_last and len(self.dataset) % self.num_replicas != 0:
            # Split to nearest available length that is evenly divisible. This
            # is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil(
                (len(self.dataset) - self.num_replicas) / self.num_replicas
            )
        else:
            self.num_samples = math.ceil(len(self.dataset) / self.num_replicas)
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def __iter__(self) -> Iterator[_T_co]:
        if self.shuffle:
            # deterministically shuffle based on epoch and seed
            g = tp.Generator()
            g.manual_seed(self.seed + self.epoch)
            with _use_generator(g):
                indices = tp.randperm(len(self.dataset)).tolist()
        else:
            indices = list(range(len(self.dataset)))

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[
                    :padding_size
                ]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[: self.total_size]
        if len(indices) != self.total_size:
            raise AssertionError(
                f"Number of indices ({len(indices)}) does not match total_size ({self.total_size})"
            )

        # subsample
        indices = indices[self.rank : self.total_size : self.num_replicas]
        if len(indices) != self.num_samples:
            raise AssertionError(
                f"Number of subsampled indices ({len(indices)}) does not match num_samples ({self.num_samples})"
            )

        return iter(indices)

    def __len__(self) -> int:
        return self.num_samples

    def set_epoch(self, epoch: int) -> None:
        r"""
        Set the epoch for this sampler.

        When :attr:`shuffle=True`, this ensures all replicas use a different
        random ordering for each epoch. Otherwise, the next iteration of this
        sampler will yield the same ordering.

        Args:
            epoch (int): Epoch number.
        """
        self.epoch = epoch

# ---------------------------------------------------------------------------
# Collation
# ---------------------------------------------------------------------------
np_str_obj_array_pattern = re.compile(r"[SaUO]")

[docs]
def default_convert(data):
    """Convert each NumPy array element into a :class:`tensorplay.Tensor`.

    If the input is a `Sequence`, `Collection`, or `Mapping`, it tries to
    convert each element inside to a :class:`tensorplay.Tensor`. If the input is
    not a NumPy array, it is left unchanged. This is used as the default
    function for collation when both `batch_sampler` and `batch_size` are NOT
    defined in :class:`~tensorplay.utils.data.DataLoader`.

    Args:
        data: a single data point to be converted
    """
    elem_type = type(data)
    if isinstance(data, Tensor):
        return data
    elif (
        elem_type.__module__ == "numpy"
        and elem_type.__name__ != "str_"
        and elem_type.__name__ != "string_"
    ):
        # array of string classes and object
        if (
            elem_type.__name__ == "ndarray"
            and np_str_obj_array_pattern.search(data.dtype.str) is not None
        ):
            return data
        return tp.as_tensor(data)
    elif isinstance(data, Mapping):
        try:
            if isinstance(data, MutableMapping):
                # The mapping type may have extra properties, so we can't just
                # use `type(data)(...)` to create the new mapping.
                # Create a clone and update it if the mapping type is mutable.
                clone = copy.copy(data)
                clone.update({key: default_convert(data[key]) for key in data})
                return clone
            else:
                return elem_type({key: default_convert(data[key]) for key in data})
        except TypeError:
            # The mapping type may not support `copy()` / `update(mapping)`
            # or `__init__(iterable)`.
            return {key: default_convert(data[key]) for key in data}
    elif isinstance(data, tuple) and hasattr(data, "_fields"):  # namedtuple
        return elem_type(*(default_convert(d) for d in data))
    elif isinstance(data, tuple):
        return [default_convert(d) for d in data]  # Backwards compatibility.
    elif isinstance(data, Sequence) and not isinstance(data, (str, bytes)):
        try:
            if isinstance(data, MutableSequence):
                # The sequence type may have extra properties, so we can't just
                # use `type(data)(...)` to create the new sequence.
                # Create a clone and update it if the sequence type is mutable.
                clone = copy.copy(data)
                for i, d in enumerate(data):
                    clone[i] = default_convert(d)
                return clone
            else:
                return elem_type([default_convert(d) for d in data])
        except TypeError:
            # The sequence type may not support `copy()` / `__setitem__(index, item)`
            # or `__init__(iterable)` (e.g., `range`).
            return [default_convert(d) for d in data]
    else:
        return data

default_collate_err_msg_format = (
    "default_collate: batch must contain tensors, numpy arrays, numbers, "
    "dicts or lists; found {}"
)

def collate(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    """General collate function that handles collection type of element within
    each batch. The function also opens function registry to deal with specific
    element types. `default_collate_fn_map` provides default collate functions
    for tensors, numpy arrays, numbers and strings.

    Args:
        batch: a single batch to be collated
        collate_fn_map: Optional dictionary mapping from element type to the
            corresponding collate function.
    """
    elem = batch[0]
    elem_type = type(elem)

    if collate_fn_map is not None:
        if elem_type in collate_fn_map:
            return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)

        for collate_type in collate_fn_map:
            if isinstance(elem, collate_type):
                return collate_fn_map[collate_type](
                    batch, collate_fn_map=collate_fn_map
                )

    if isinstance(elem, Mapping):
        try:
            if isinstance(elem, MutableMapping):
                # The mapping type may have extra properties, so we can't just
                # use `type(data)(...)` to create the new mapping.
                # Create a clone and update it if the mapping type is mutable.
                clone = copy.copy(elem)
                clone.update(
                    {
                        key: collate(
                            [d[key] for d in batch], collate_fn_map=collate_fn_map
                        )
                        for key in elem
                    }
                )
                return clone
            else:
                return elem_type(
                    {
                        key: collate(
                            [d[key] for d in batch], collate_fn_map=collate_fn_map
                        )
                        for key in elem
                    }
                )
        except TypeError:
            # The mapping type may not support `copy()` / `update(mapping)`
            # or `__init__(iterable)`.
            return {
                key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map)
                for key in elem
            }
    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
        return elem_type(
            *(
                collate(samples, collate_fn_map=collate_fn_map)
                for samples in zip(*batch, strict=False)
            )
        )
    elif isinstance(elem, Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))

        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError("each element in list of batch should be of equal size")
        transposed = list(
            zip(*batch, strict=False)
        )  # It may be accessed twice, so we use a list.

        if isinstance(elem, tuple):
            return [
                collate(samples, collate_fn_map=collate_fn_map)
                for samples in transposed
            ]  # Backwards compatibility.
        else:
            try:
                if isinstance(elem, MutableSequence):
                    # The sequence type may have extra properties, so we can't just
                    # use `type(data)(...)` to create the new sequence.
                    # Create a clone and update it if the sequence type is mutable.
                    clone = copy.copy(elem)
                    for i, samples in enumerate(transposed):
                        clone[i] = collate(samples, collate_fn_map=collate_fn_map)
                    return clone
                else:
                    return elem_type(
                        [
                            collate(samples, collate_fn_map=collate_fn_map)
                            for samples in transposed
                        ]
                    )
            except TypeError:
                # The sequence type may not support `copy()` / `__setitem__(index, item)`
                # or `__init__(iterable)` (e.g., `range`).
                return [
                    collate(samples, collate_fn_map=collate_fn_map)
                    for samples in transposed
                ]

    raise TypeError(default_collate_err_msg_format.format(elem_type))

def collate_tensor_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    return tp.stack(batch, 0)

def collate_numpy_array_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    elem = batch[0]
    # array of string classes and object
    if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
        raise TypeError(default_collate_err_msg_format.format(elem.dtype))

    return collate([tp.as_tensor(b) for b in batch], collate_fn_map=collate_fn_map)

def collate_numpy_scalar_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    return tp.as_tensor(np.asarray(batch))

def collate_float_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    return tp.tensor(np.asarray(batch), dtype=tp.float64)

def collate_int_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    return tp.tensor(np.asarray(batch))

def collate_str_fn(
    batch,
    *,
    collate_fn_map: Optional[dict] = None,
):
    return batch

default_collate_fn_map: dict = {
    Tensor: collate_tensor_fn,
    np.ndarray: collate_numpy_array_fn,
    (np.bool_, np.number, np.object_): collate_numpy_scalar_fn,
    float: collate_float_fn,
    int: collate_int_fn,
    str: collate_str_fn,
    bytes: collate_str_fn,
}

[docs]
def default_collate(batch):
    """Take in a batch of data and put the elements within the batch into a
    tensor with an additional outer dimension - batch size.

    The exact output type can be a :class:`tensorplay.Tensor`, a `Sequence` of
    :class:`tensorplay.Tensor`, a Collection of :class:`tensorplay.Tensor`, or
    left unchanged, depending on the input type. This is used as the default
    function for collation when `batch_size` or `batch_sampler` is defined in
    :class:`~tensorplay.utils.data.DataLoader`.

    Args:
        batch: a single batch to be collated
    """
    return collate(batch, collate_fn_map=default_collate_fn_map)

# ---------------------------------------------------------------------------
# Worker helpers (multiprocessing)
# ---------------------------------------------------------------------------
class _DatasetKind:
    Map = 0
    Iterable = 1

    @staticmethod
    def create_fetcher(dataset_kind, dataset, auto_collation, collate_fn, drop_last):
        if dataset_kind == _DatasetKind.Iterable:
            return _IterableDatasetFetcher(dataset, auto_collation, collate_fn, drop_last)
        else:
            return _MapDatasetFetcher(dataset, auto_collation, collate_fn, drop_last)

_MP_STATUS_CHECK_INTERVAL = 5.0

_worker_info: Optional["WorkerInfo"] = None

class WorkerInfo:
    """Information about the current DataLoader worker process.

    Attributes:
        id: The current worker id (0 to num_workers - 1).
        num_workers: Total number of workers.
        seed: Random seed set for this worker (``base_seed + worker_id``).
        dataset: Copy of the dataset object in this worker.
    """

    __slots__ = ("id", "num_workers", "seed", "dataset")

    def __init__(self, id, num_workers, seed, dataset):
        self.id = id
        self.num_workers = num_workers
        self.seed = seed
        self.dataset = dataset

    def __repr__(self):
        return (
            f"<WorkerInfo(id={self.id}, num_workers={self.num_workers}, "
            f"seed={self.seed}, dataset={self.dataset!r})>"
        )

[docs]
def get_worker_info() -> Optional[WorkerInfo]:
    """Returns the information about the current DataLoader iterator worker
    process. When called in a worker process, returns a WorkerInfo object with
    information about that worker process; otherwise returns ``None``.
    """
    return _worker_info

class _KeyErrorMessage(str):
    # repr() of a str is surrounded by quotes, which is unreadable for KeyError
    def __repr__(self):
        return self

class ExceptionWrapper:
    """Wraps an exception plus traceback to communicate across processes."""

    def __init__(self, exc_info=None, where="in background"):
        # It is important that we don't store exc_info in a variable.
        if exc_info is None:
            exc_info = sys.exc_info()
        self.exc_type = exc_info[0]
        self.exc_msg = "".join(traceback.format_exception(*exc_info))
        self.where = where

    def reraise(self):
        msg = f"Caught {self.exc_type.__name__} {self.where}.\nOriginal {self.exc_msg}"
        if self.exc_type is KeyError:
            msg = _KeyErrorMessage(msg)
        elif getattr(self.exc_type, "message", None):
            raise self.exc_type(message=msg)
        try:
            exception = self.exc_type(msg)
        except Exception:
            raise RuntimeError(msg) from None
        raise exception

# ---------------------------------------------------------------------------
# Dataset fetchers (shared by single- and multi-process loading)
# ---------------------------------------------------------------------------
class _BaseDatasetFetcher:
    def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None:
        self.dataset = dataset
        self.auto_collation = auto_collation
        self.collate_fn = collate_fn
        self.drop_last = drop_last

    def fetch(self, possibly_batched_index):
        raise NotImplementedError

class _IterableDatasetFetcher(_BaseDatasetFetcher):
    def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None:
        super().__init__(dataset, auto_collation, collate_fn, drop_last)
        self.dataset_iter = iter(dataset)
        self.ended = False

    def fetch(self, possibly_batched_index):
        if self.ended:
            raise StopIteration

        if self.auto_collation:
            data = []
            for _ in possibly_batched_index:
                try:
                    data.append(next(self.dataset_iter))
                except StopIteration:
                    self.ended = True
                    break
            if len(data) == 0 or (
                self.drop_last and len(data) < len(possibly_batched_index)
            ):
                raise StopIteration
        else:
            data = next(self.dataset_iter)
        return self.collate_fn(data)

class _MapDatasetFetcher(_BaseDatasetFetcher):
    def fetch(self, possibly_batched_index):
        if self.auto_collation:
            # TensorPlay-backed datasets can return the already-collated
            # result of a batch index_select.  This is deliberately gated on
            # the stock collate function: a user collate_fn must still receive
            fast_getter = getattr(self.dataset, "_tp_get_batch", None)
            if self.collate_fn is default_collate and callable(fast_getter):
                data = fast_getter(possibly_batched_index)
                if data is not _FAST_BATCH_UNAVAILABLE:
                    return data

            getitems = getattr(self.dataset, "__getitems__", None)
            if callable(getitems):
                data = getitems(possibly_batched_index)
            else:
                data = [self.dataset[idx] for idx in possibly_batched_index]
        else:
            data = self.dataset[possibly_batched_index]
        return self.collate_fn(data)

# The function `_generate_state` is adapted from `numpy.random.SeedSequence`
# It generates an array of int32 as the seed for `numpy.random`, in order to
# prevent state collision due to same seed and algorithm for `numpy.random`
# and `random` modules.
def _generate_state(base_seed, worker_id):
    INIT_A = 0x43B0D7E5
    MULT_A = 0x931E8875
    INIT_B = 0x8B51F9DD
    MULT_B = 0x58F38DED
    MIX_MULT_L = 0xCA01F9DD
    MIX_MULT_R = 0x4973F715
    XSHIFT = 4 * 8 // 2
    MASK32 = 0xFFFFFFFF

    entropy = [worker_id, base_seed & MASK32, base_seed >> 32, 0]
    pool = [0] * 4

    hash_const_A = INIT_A

    def hash(value):
        nonlocal hash_const_A
        value = (value ^ hash_const_A) & MASK32
        hash_const_A = (hash_const_A * MULT_A) & MASK32
        value = (value * hash_const_A) & MASK32
        value = (value ^ (value >> XSHIFT)) & MASK32
        return value

    def mix(x, y):
        result_x = (MIX_MULT_L * x) & MASK32
        result_y = (MIX_MULT_R * y) & MASK32
        result = (result_x - result_y) & MASK32
        result = (result ^ (result >> XSHIFT)) & MASK32
        return result

    for i in range(len(pool)):
        pool[i] = hash(entropy[i])

    for i_src in range(len(pool)):
        for i_dst in range(len(pool)):
            if i_src != i_dst:
                pool[i_dst] = mix(pool[i_dst], hash(pool[i_src]))

    hash_const_B = INIT_B
    state = []
    for i_dst in range(4):
        data_val = pool[i_dst]
        data_val = (data_val ^ hash_const_B) & MASK32
        hash_const_B = (hash_const_B * MULT_B) & MASK32
        data_val = (data_val * hash_const_B) & MASK32
        data_val = (data_val ^ (data_val >> XSHIFT)) & MASK32
        state.append(data_val)
    return state

class _IterableDatasetStopIteration:
    """Dummy class used to signal the end of an IterableDataset worker."""

    def __init__(self, worker_id):
        self.worker_id = worker_id

class _ResumeIteration:
    """Dummy class used to resume fetching when worker reuse is enabled."""

    def __init__(self, seed=None):
        self.seed = seed

class _ManagerWatchdog:
    def __init__(self) -> None:
        self.manager_pid = os.getppid()
        self.manager_dead = False

    def is_alive(self) -> bool:
        if not self.manager_dead:
            self.manager_dead = os.getppid() != self.manager_pid
        return not self.manager_dead

def _worker_loop(
    dataset_kind,
    dataset,
    index_queue,
    data_queue,
    done_event,
    auto_collation,
    collate_fn,
    drop_last,
    base_seed,
    init_fn,
    worker_id,
    num_workers,
    persistent_workers,
) -> None:
    # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] in
    try:
        watchdog = _ManagerWatchdog()

        tp.set_num_threads(1)
        seed = base_seed + worker_id
        random.seed(seed)
        tp.manual_seed(seed)
        np.random.seed(_generate_state(base_seed, worker_id))

        global _worker_info
        _worker_info = WorkerInfo(
            id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset
        )

        init_exception = None

        try:
            if init_fn is not None:
                init_fn(worker_id)

            fetcher = _DatasetKind.create_fetcher(
                dataset_kind, dataset, auto_collation, collate_fn, drop_last
            )
        except Exception:
            init_exception = ExceptionWrapper(
                where=f"in DataLoader worker process {worker_id}"
            )

        # When using Iterable mode, some workers can exit earlier than others
        # due to the IterableDataset behaving differently per worker. Such a
        # worker sends `_IterableDatasetStopIteration` so the main process
        # won't send more tasks to it and will send `None` to exit it.
        iteration_end = False

        while watchdog.is_alive():
            try:
                r = index_queue.get(timeout=_MP_STATUS_CHECK_INTERVAL)
            except queue.Empty:
                continue
            if isinstance(r, _ResumeIteration):
                # Acknowledge the main process
                data_queue.put((r, None))
                iteration_end = False
                # Recreate the fetcher for the worker-reuse policy
                fetcher = _DatasetKind.create_fetcher(
                    dataset_kind, dataset, auto_collation, collate_fn, drop_last
                )
                continue
            elif r is None:
                # Received the final signal
                break
            elif done_event.is_set() or iteration_end:
                # Keep getting until we see the final signal (`None`), skipping
                # the processing steps.
                continue
            idx, index = r
            if init_exception is not None:
                data = init_exception
                init_exception = None
            else:
                try:
                    data = fetcher.fetch(index)
                except Exception as e:
                    if (
                        isinstance(e, StopIteration)
                        and dataset_kind == _DatasetKind.Iterable
                    ):
                        data = _IterableDatasetStopIteration(worker_id)
                        iteration_end = True
                    else:
                        data = ExceptionWrapper(
                            where=f"in DataLoader worker process {worker_id}"
                        )
            data_queue.put((idx, data))
    except KeyboardInterrupt:
        # Main process will raise KeyboardInterrupt anyways.
        pass
    finally:
        _worker_info = None
    if done_event.is_set():
        data_queue.cancel_join_thread()
        data_queue.close()

# ---------------------------------------------------------------------------
# Pin memory
# ---------------------------------------------------------------------------
def _pin_memory(data):
    if isinstance(data, Tensor):
        return data.pin_memory()

    if hasattr(data, "pin_memory"):
        return data.pin_memory()

    if isinstance(data, (str, bytes)):
        return data

    if isinstance(data, Mapping):
        try:
            if isinstance(data, MutableMapping):
                clone = copy.copy(data)
                clone.update({k: _pin_memory(sample) for k, sample in data.items()})
                return clone
            else:
                return type(data)({k: _pin_memory(sample) for k, sample in data.items()})
        except TypeError:
            return {k: _pin_memory(sample) for k, sample in data.items()}

    if isinstance(data, tuple):
        if hasattr(data, "_fields"):  # namedtuple
            return type(data)(*(_pin_memory(sample) for sample in data))
        return type(data)(_pin_memory(sample) for sample in data)

    if isinstance(data, Sequence):
        try:
            if isinstance(data, MutableSequence):
                clone = copy.copy(data)
                for i, item in enumerate(data):
                    clone[i] = _pin_memory(item)
                return clone
            return type(data)([_pin_memory(sample) for sample in data])
        except TypeError:
            return [_pin_memory(sample) for sample in data]

    return data

def _pin_memory_loop(in_queue, out_queue, done_event) -> None:
    def do_one_step():
        try:
            r = in_queue.get(timeout=_MP_STATUS_CHECK_INTERVAL)
        except queue.Empty:
            return
        idx, data = r
        if not done_event.is_set() and not isinstance(data, ExceptionWrapper):
            try:
                data = _pin_memory(data)
            except Exception:
                data = ExceptionWrapper(where="in pin memory thread")
            r = (idx, data)
        while not done_event.is_set():
            try:
                out_queue.put(r, timeout=_MP_STATUS_CHECK_INTERVAL)
                break
            except queue.Full:
                continue

    while not done_event.is_set():
        do_one_step()

def _default_multiprocessing_context():
    try:
        methods = multiprocessing.get_all_start_methods()
    except Exception:
        methods = []
    if "fork" in methods:
        return multiprocessing.get_context("fork")
    return multiprocessing.get_context()

def _new_base_seed(generator: Optional["tp.Generator"]) -> int:
    with _use_generator(generator):
        return int(tp.empty((), dtype=tp.int64).random_(0, 2**63 - 1).item())

def _cpu_count() -> Optional[int]:
    """

    Prefers ``os.sched_getaffinity`` (respects cgroups / taskset) and falls
    back to ``os.cpu_count``.
    """
    if hasattr(os, "sched_getaffinity"):
        return len(os.sched_getaffinity(0))
    return os.cpu_count()

def _worker_rationality_warning_msg(num_worker_suggest, num_worker_created,
                                    cpuset_checked: bool) -> str:
    suggested_max_worker_msg = (
        (
            (
                "Our suggested max number of worker in current system is {}{}, which is smaller "
                "than what this DataLoader is going to create."
            ).format(
                num_worker_suggest,
                "" if cpuset_checked else " (`cpuset` is not taken into account)",
            )
        )
        if num_worker_suggest is not None
        else (
            "DataLoader is not able to compute a suggested max number of worker in current system."
        )
    )
    return (
        f"This DataLoader will create {num_worker_created} worker processes in total. {suggested_max_worker_msg} "
        "Please be aware that excessive worker creation might get DataLoader running slow or even freeze, "
        "lower the worker number to avoid potential slowness/freeze if necessary."
    )

class _InfiniteConstantSampler(Sampler[int]):
    r"""Analogous to ``itertools.repeat(0, None)``, but with a ``__len__`` that
    raises ``TypeError``."""

    def __iter__(self) -> Iterator[int]:
        while True:
            yield 0

    def __len__(self) -> int:
        raise TypeError("Cannot determine the length of an infinite sampler")

# ---------------------------------------------------------------------------
# DataLoader
# ---------------------------------------------------------------------------
class _BaseDataLoaderIter:
    def __init__(self, loader: "DataLoader"):
        self.loader = loader
        self._dataset_kind = loader._dataset_kind
        self._IterableDataset_len_called = loader._IterableDataset_len_called
        self._auto_collation = loader._auto_collation
        self._drop_last = loader.drop_last
        self._index_sampler = loader._index_sampler
        self._sampler_iter = iter(self._index_sampler)
        self._num_workers = loader.num_workers
        self._num_yielded = 0

        # memory is enabled only when an accelerator exists.  An explicit
        # device opts into the requested allocator and therefore keeps the
        # flag even when the current process has no default accelerator.
        if not loader.pin_memory_device:
            if loader.pin_memory and not tp.cuda.is_available():
                warnings.warn(
                    "'pin_memory' argument is set as true but no accelerator is found, "
                    "then device pinned memory won't be used.",
                    stacklevel=2,
                )
            self._pin_memory = bool(loader.pin_memory and tp.cuda.is_available())
            self._pin_memory_device = None
        else:
            if not loader.pin_memory:
                warnings.warn(
                    "'pin_memory_device' is set but 'pin_memory' argument is not set, "
                    "then device pinned memory won't be used."
                    "please set 'pin_memory' to true, if you need to use the device pin memory",
                    stacklevel=2,
                )
            self._pin_memory = bool(loader.pin_memory)
            self._pin_memory_device = loader.pin_memory_device
        self._base_seed = _new_base_seed(loader.generator)

    def __iter__(self) -> "_BaseDataLoaderIter":
        return self

    def _next_data(self):
        raise NotImplementedError

    def __next__(self) -> Any:
        data = self._next_data()
        self._num_yielded += 1
        if (
            self._dataset_kind == _DatasetKind.Iterable
            and self._IterableDataset_len_called is not None
            and self._num_yielded > self._IterableDataset_len_called
        ):
            # See NOTE [ IterableDataset and __len__ ] in ``DataLoader.__len__``.
            warn_msg = (
                f"Length of IterableDataset {self.loader.dataset} was reported to be {self._IterableDataset_len_called}"
                f"(when accessing len(dataloader)), but {self._num_yielded} samples have been fetched. "
            )
            if self.loader.num_workers > 0:
                warn_msg += (
                    "For multiprocessing data-loading, this could be caused by not properly configuring the "
                    "IterableDataset replica at each worker. Please see "
                )
            warnings.warn(warn_msg, stacklevel=2)
        return data

class _SingleProcessDataLoaderIter(_BaseDataLoaderIter):
    def __init__(self, loader: "DataLoader"):
        super().__init__(loader)
        self._dataset_fetcher = _DatasetKind.create_fetcher(
            loader._dataset_kind,
            loader.dataset,
            loader._auto_collation,
            loader.collate_fn,
            loader.drop_last,
        )

    def _next_data(self) -> Any:
        index = next(self._sampler_iter)  # may raise StopIteration
        data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
        return self.loader._finalize_batch(data)

class _MultiProcessDataLoaderIter(_BaseDataLoaderIter):
    r"""Iterates once over the DataLoader's dataset, as specified by the sampler.

    each worker runs fetch + collate on a per-worker index queue, results are
    streamed back through a shared result queue with bounded prefetch, and an
    optional pin-memory thread of the main process pins the returned batches.
    """

    # NOTE [ Data Loader Multiprocessing Shutdown Logic ]
    #
    # Workers are daemonic. The protocol between the main process and the
    # workers is that the main process sets `workers_done_event` (or marks a
    # worker unavailable for IterableDataset exhaustion), then sends `None` on
    # the worker's index queue; workers exit upon receiving `None`. Normal
    # data always travels as a 2-tuple `(idx, data)` where `idx` is the task
    # index assigned by `_try_put_index`, so `None` can never be confused with
    # the full note.

    def __init__(self, loader: "DataLoader"):
        super().__init__(loader)
        self._timeout = loader.timeout
        self._collate_fn = loader.collate_fn
        self._num_workers = loader.num_workers
        self._prefetch_factor = loader.prefetch_factor
        self._persistent_workers = loader.persistent_workers
        self._in_order = loader.in_order

        if self._num_workers <= 0:
            raise AssertionError(
                "num_workers must be greater than 0 for MultiProcessingDataLoaderIter"
            )
        if self._prefetch_factor <= 0:
            raise AssertionError(
                "prefetch_factor must be greater than 0 for MultiProcessingDataLoaderIter"
            )

        if loader.multiprocessing_context is None:
            multiprocessing_context = _default_multiprocessing_context()
        else:
            multiprocessing_context = loader.multiprocessing_context

        self._worker_init_fn = loader.worker_init_fn

        self._worker_result_queue = multiprocessing_context.Queue()
        self._shutdown = False
        self._workers_done_event = multiprocessing_context.Event()

        self._index_queues = []
        self._workers = []
        for i in range(self._num_workers):
            index_queue = multiprocessing_context.Queue()
            # Need to `cancel_join_thread` here! See sections (2) and (3b) of
            # NOTE [ Data Loader Multiprocessing Shutdown Logic ].
            index_queue.cancel_join_thread()
            w = multiprocessing_context.Process(
                target=_worker_loop,
                args=(
                    self._dataset_kind,
                    loader.dataset,
                    index_queue,
                    self._worker_result_queue,
                    self._workers_done_event,
                    self._auto_collation,
                    self._collate_fn,
                    self._drop_last,
                    self._base_seed,
                    self._worker_init_fn,
                    i,
                    self._num_workers,
                    self._persistent_workers,
                ),
            )
            w.daemon = True
            w.start()
            self._index_queues.append(index_queue)
            self._workers.append(w)

        if self._pin_memory:
            self._pin_memory_thread_done_event = threading.Event()

            # The pin thread consumes the worker result queue and feeds a
            # plain `queue.Queue` that `__next__` reads.
            self._data_queue: Any = queue.Queue()
            pin_memory_thread = threading.Thread(
                target=_pin_memory_loop,
                args=(
                    self._worker_result_queue,
                    self._data_queue,
                    self._pin_memory_thread_done_event,
                ),
            )
            pin_memory_thread.daemon = True
            pin_memory_thread.start()
            self._pin_memory_thread = pin_memory_thread
        else:
            self._data_queue = self._worker_result_queue

        self._reset(loader, first_iter=True)

    def _reset(self, loader: "DataLoader", first_iter: bool = False) -> None:
        self._index_sampler = loader._index_sampler
        self._sampler_iter = iter(self._index_sampler)
        self._IterableDataset_len_called = loader._IterableDataset_len_called
        self._send_idx = 0  # idx of the next task to be sent to workers
        self._rcvd_idx = 0  # idx of the next task to be returned in __next__
        # map: task idx => (worker_id,)     if data isn't fetched (outstanding)
        #                  \ (worker_id, data) if already fetched (out-of-order)
        self._task_info = {}
        self._tasks_outstanding = 0
        # Whether each worker still has work to do *for this epoch* (i.e., has
        # not exhausted its IterableDataset replica). All `True`s for map-style
        # datasets.
        self._workers_status = [True for i in range(self._num_workers)]
        # Outstanding-task count per worker; each worker holds at most
        # `prefetch_factor` tasks when `in_order=False` load-balances dispatch.
        self._workers_num_tasks = [0 for i in range(self._num_workers)]
        # Reset the worker queue cycle so a new epoch resumes at worker 0.
        self._worker_queue_idx_cycle = itertools.cycle(range(self._num_workers))
        # We resume the prefetching in case it was enabled (persistent workers).
        if not first_iter:
            for idx in range(self._num_workers):
                self._index_queues[idx].put(_ResumeIteration())
            resume_iteration_cnt = self._num_workers
            while resume_iteration_cnt > 0:
                return_idx, return_data = self._get_data()
                if isinstance(return_idx, _ResumeIteration):
                    if return_data is not None:
                        raise AssertionError(
                            "Expected return_data to be None when resuming iteration"
                        )
                    resume_iteration_cnt -= 1
        # prime the prefetch loop
        for _ in range(self._prefetch_factor * self._num_workers):
            self._try_put_index()

    def _try_get_data(self, timeout=_MP_STATUS_CHECK_INTERVAL):
        # Tries to fetch data from `self._data_queue` once for a given timeout.
        # Raises RuntimeError if any worker died unexpectedly. Returns a
        # 2-tuple (success, data).
        try:
            data = self._data_queue.get(timeout=timeout)
            return (True, data)
        except Exception as e:
            # At timeout and error, we manually check whether any worker has
            # failed.
            failed_workers = []
            for worker_id, w in enumerate(self._workers):
                if self._workers_status[worker_id] and not w.is_alive():
                    failed_workers.append(w)
                    self._mark_worker_as_unavailable(worker_id)
            if len(failed_workers) > 0:
                pids_str = ", ".join(str(w.pid) for w in failed_workers)
                raise RuntimeError(
                    f"DataLoader worker (pid(s) {pids_str}) exited unexpectedly"
                ) from e
            if isinstance(e, queue.Empty):
                return (False, None)
            raise

    def _get_data(self):
        # Fetches data from `self._data_queue`, checking worker/pin-thread
        # status every `_MP_STATUS_CHECK_INTERVAL` seconds.
        if self._timeout > 0:
            success, data = self._try_get_data(self._timeout)
            if success:
                return data
            else:
                raise RuntimeError(
                    f"DataLoader timed out after {self._timeout} seconds"
                )
        elif hasattr(self, "_pin_memory_thread"):
            while self._pin_memory_thread.is_alive():
                success, data = self._try_get_data()
                if success:
                    return data
            # while condition is false, i.e., pin_memory_thread died.
            raise RuntimeError("Pin memory thread exited unexpectedly")
        else:
            while True:
                success, data = self._try_get_data()
                if success:
                    return data

    def _next_data(self):
        while True:
            # If the worker responsible for `self._rcvd_idx` has already ended
            # and was unable to fulfill this task (due to exhausting an
            # IterableDataset), advance `self._rcvd_idx` to find the next
            # valid index.
            while self._rcvd_idx < self._send_idx:
                info = self._task_info.get(self._rcvd_idx, None)
                if info:
                    worker_id = info[0]
                    if len(info) == 2 or self._workers_status[worker_id]:
                        break
                    del self._task_info[self._rcvd_idx]
                self._rcvd_idx += 1
            else:
                # no valid `self._rcvd_idx` is found (i.e., didn't break)
                if not self._persistent_workers:
                    self._shutdown_workers()
                raise StopIteration

            # Now `self._rcvd_idx` is the batch index we want to fetch

            # Check if the next sample has already been generated
            if len(self._task_info[self._rcvd_idx]) == 2:
                worker_id, data = self._task_info.pop(self._rcvd_idx)
                self._rcvd_idx += 1
                return self._process_data(data, worker_id)

            if self._shutdown or self._tasks_outstanding <= 0:
                raise AssertionError(
                    "Invalid iterator state: shutdown or no outstanding tasks when fetching next data"
                )
            idx, data = self._get_data()
            self._tasks_outstanding -= 1
            if self._dataset_kind == _DatasetKind.Iterable:
                # Check for _IterableDatasetStopIteration
                if isinstance(data, _IterableDatasetStopIteration):
                    if self._persistent_workers:
                        self._workers_status[data.worker_id] = False
                    else:
                        self._mark_worker_as_unavailable(data.worker_id)
                    self._try_put_index()
                    continue

            if idx != self._rcvd_idx:
                if not self._in_order:
                    # don't store it for later, process now; this keeps the
                    # object size manageable
                    worker_id = self._task_info.pop(idx)[0]
                    return self._process_data(data, worker_id)
                # store out-of-order samples
                self._task_info[idx] += (data,)
            else:
                worker_id = self._task_info.pop(idx)[0]
                self._rcvd_idx += 1
                return self._process_data(data, worker_id)

    def _try_put_index(self) -> None:
        max_tasks = self._prefetch_factor * self._num_workers
        if self._tasks_outstanding >= max_tasks:
            raise AssertionError(
                "Number of outstanding tasks exceeded maximum allowed tasks"
            )

        try:
            index = next(self._sampler_iter)
        except StopIteration:
            return
        for _ in range(self._num_workers):  # find the next active worker, if any
            worker_queue_idx = next(self._worker_queue_idx_cycle)
            if self._workers_status[worker_queue_idx]:
                if self._in_order:
                    break
                elif self._workers_num_tasks[worker_queue_idx] < max_tasks // sum(
                    self._workers_status
                ):
                    # when in_order is False, distribute work to a worker with
                    # capacity; _workers_status only changes in this thread, so
                    # the sum is guaranteed > 0.
                    break
        else:
            # not found (i.e., didn't break)
            return

        self._index_queues[worker_queue_idx].put((self._send_idx, index))
        self._task_info[self._send_idx] = (worker_queue_idx,)
        self._workers_num_tasks[worker_queue_idx] += 1
        self._tasks_outstanding += 1
        self._send_idx += 1

    def _process_data(self, data, worker_idx):
        self._workers_num_tasks[worker_idx] -= 1
        self._try_put_index()
        if isinstance(data, ExceptionWrapper):
            data.reraise()
        # Pinning (if enabled) already happened in the pin-memory thread;
        # the optional device move is applied here in the main process.
        return self.loader._apply_device(data)

    def _mark_worker_as_unavailable(self, worker_id, shutdown=False) -> None:
        # Mark a worker as having finished its work e.g., due to exhausting an
        # IterableDataset. Signal termination to that specific worker; joining
        # is deferred to `_shutdown_workers`.
        q = self._index_queues[worker_id]
        q.put(None)
        self._workers_status[worker_id] = False

    def _shutdown_workers(self) -> None:
        # Exit the pin memory thread first because exiting workers may leave
        # corrupted data in `worker_result_queue` which it reads from; then
        # signal and join the workers.
        if self._shutdown:
            return
        self._shutdown = True
        try:
            if hasattr(self, "_pin_memory_thread"):
                self._pin_memory_thread_done_event.set()
                # Send something to wake the pin thread up in case it waits.
                self._worker_result_queue.put((None, None))
                self._pin_memory_thread.join()
                self._worker_result_queue.cancel_join_thread()
                self._worker_result_queue.close()

            self._workers_done_event.set()
            for worker_id in range(len(self._workers)):
                if self._persistent_workers or self._workers_status[worker_id]:
                    self._mark_worker_as_unavailable(worker_id, shutdown=True)
            for w in self._workers:
                w.join(timeout=_MP_STATUS_CHECK_INTERVAL)
            for q in self._index_queues:
                q.cancel_join_thread()
                q.close()
        finally:
            for w in self._workers:
                if w.is_alive():
                    # Existing mechanisms try to make the workers exit
                    # peacefully, but in case we unfortunately reach here, we
                    # kill the worker.
                    w.terminate()
                    w.join(timeout=_MP_STATUS_CHECK_INTERVAL)

    def __del__(self):
        try:
            self._shutdown_workers()
        except Exception:
            pass

[docs]
class DataLoader(Generic[_T_co]):
    """Data loader combines a dataset and a sampler, and provides an iterable
    over the given dataset.

    The DataLoader supports both map-style and iterable-style datasets with
    single- or multi-process loading, customizing loading order and optional
    automatic batching (collation).

    Args:
        dataset (Dataset): dataset from which to load the data.
        batch_size (int, optional): how many samples per batch to load
            (default: ``1``).
        shuffle (bool, optional): set to ``True`` to have the data reshuffled
            at every epoch (default: ``False``).
        sampler (Sampler or Iterable, optional): defines the strategy to draw
            samples from the dataset. Can be any ``Iterable`` with ``__len__``
            implemented. If specified, :attr:`shuffle` must not be specified.
        batch_sampler (Sampler or Iterable, optional): like :attr:`sampler`, but
            returns a batch of indices at a time. Mutually exclusive with
            :attr:`batch_size`, :attr:`shuffle`, :attr:`sampler`, and
            :attr:`drop_last`.
        num_workers (int, optional): how many subprocesses to use for data
            loading. ``0`` means that the data will be loaded in the main
            process. (default: ``0``)
        collate_fn (Callable, optional): merges a list of samples to form a
            mini-batch of Tensor(s). Used when using batched loading from a
            map-style dataset.
        pin_memory (bool, optional): If ``True``, the data loader will copy
            Tensors into CUDA page-locked host memory before returning them.
        drop_last (bool, optional): set to ``True`` to drop the last incomplete
            batch, if the dataset size is not divisible by the batch size.
            (default: ``False``)
        timeout (numeric, optional): if positive, the timeout value for
            collecting a batch from workers. Should always be non-negative.
            (default: ``0``)
        worker_init_fn (Callable, optional): If not ``None``, this will be
            called on each worker subprocess with the worker id (an int in
            ``[0, num_workers - 1]``) as input, before data loading.
            (default: ``None``)
        multiprocessing_context (str or context, optional): start method or
            ``multiprocessing`` context used to spawn the workers, e.g.,
            ``"fork"`` or ``"spawn"``. If ``None``, the default context of the
            platform is used. (default: ``None``)
        generator (Generator, optional): If not ``None``, this RNG will be used
            by RandomSampler to generate random indexes. (default: ``None``)
        prefetch_factor (int, optional): Number of batches loaded in advance by
            each worker. ``2`` means there will be a total of
            2 * :attr:`num_workers` batches prefetched across all workers.
            (default: ``2`` when :attr:`num_workers > 0`; otherwise must be
            ``None``)
        persistent_workers (bool, optional): If ``True``, the data loader will
            not shut down the worker processes after a dataset has been
            consumed once. This allows to maintain the workers `Dataset`
            instances alive. (default: ``False``)
        pin_memory_device (str, optional): Deprecated device spelling kept for
            accelerator for pinned host allocations. (default: ``""``)
        in_order (bool, optional): If ``False``, the data loader will not
            enforce that batches returned from multiprocessing workers are
            provided in the order the sampler produced them. This enables
            faster delivery of batches that complete early, at the cost of
            batch order no longer being deterministic. (default: ``True``)
        device (str, optional): device to move batches to after collation.
    """

    dataset: Dataset[_T_co]
    batch_size: Optional[int]
    num_workers: int
    pin_memory: bool
    pin_memory_device: str
    drop_last: bool
    timeout: float
    sampler: Sampler | Iterable
    generator: Optional["tp.Generator"]
    __initialized = False

    def __init__(
        self,
        dataset: Dataset[_T_co],
        batch_size: Optional[int] = 1,
        shuffle: Optional[bool] = None,
        sampler: Optional[Sampler | Iterable] = None,
        batch_sampler: Optional[Sampler | Iterable] = None,
        num_workers: int = 0,
        collate_fn: Optional[Callable[[List[Any]], Any]] = None,
        pin_memory: bool = False,
        drop_last: bool = False,
        timeout: float = 0,
        worker_init_fn: Optional[Callable[[int], None]] = None,
        multiprocessing_context=None,
        generator: Optional["tp.Generator"] = None,
        *,
        prefetch_factor: Optional[int] = None,
        persistent_workers: bool = False,
        pin_memory_device: str = "",
        in_order: bool = True,
        device: Optional[str] = None,
    ) -> None:
        if num_workers < 0:
            raise ValueError(
                "num_workers option should be non-negative; "
                "use num_workers=0 to disable multiprocessing."
            )

        if timeout < 0:
            raise ValueError("timeout option should be non-negative")

        if num_workers == 0 and prefetch_factor is not None:
            raise ValueError(
                "prefetch_factor option could only be specified in multiprocessing."
                "let num_workers > 0 to enable multiprocessing, otherwise set prefetch_factor to None."
            )
        elif num_workers > 0 and prefetch_factor is None:
            prefetch_factor = 2
        elif prefetch_factor is not None and prefetch_factor < 0:
            raise ValueError("prefetch_factor option should be non-negative")

        if persistent_workers and num_workers == 0:
            raise ValueError("persistent_workers option needs num_workers > 0")

        self.dataset = dataset
        self.num_workers = num_workers
        self.prefetch_factor = prefetch_factor
        self.pin_memory = pin_memory
        self.pin_memory_device = pin_memory_device
        self.timeout = timeout
        self.worker_init_fn = worker_init_fn
        self.multiprocessing_context = multiprocessing_context
        self.persistent_workers = persistent_workers
        self.in_order = in_order
        self.generator = generator
        self.device = device

        # Arg-check dataset related before checking samplers because we want to
        # tell users that iterable-style datasets are incompatible with custom
        # samplers first, so that they don't learn that this combo doesn't work
        # after spending time fixing the custom sampler errors.
        if isinstance(dataset, IterableDataset):
            self._dataset_kind = _DatasetKind.Iterable
            # We cannot check `shuffle is not None` here, since previously
            # `shuffle=False` was the default.
            if shuffle not in {False, None}:
                raise ValueError(
                    f"DataLoader with IterableDataset: expected unspecified shuffle option, but got shuffle={shuffle}"
                )
            if sampler is not None:
                raise ValueError(
                    f"DataLoader with IterableDataset: expected unspecified sampler option, but got sampler={sampler}"
                )
            elif batch_sampler is not None:
                raise ValueError(
                    "DataLoader with IterableDataset: expected unspecified "
                    f"batch_sampler option, but got batch_sampler={batch_sampler}"
                )
        else:
            shuffle = bool(shuffle)
            self._dataset_kind = _DatasetKind.Map

        if sampler is not None and shuffle:
            raise ValueError("sampler option is mutually exclusive with shuffle")

        if batch_sampler is not None:
            # auto_collation with custom batch_sampler
            if batch_size != 1 or shuffle or sampler is not None or drop_last:
                raise ValueError(
                    "batch_sampler option is mutually exclusive "
                    "with batch_size, shuffle, sampler, and "
                    "drop_last"
                )
            batch_size = None
            drop_last = False
        elif batch_size is None:
            # no auto_collation
            if drop_last:
                raise ValueError(
                    "batch_size=None option disables auto-batching "
                    "and is mutually exclusive with drop_last"
                )

        if sampler is None:  # give default samplers
            if self._dataset_kind == _DatasetKind.Iterable:
                sampler = _InfiniteConstantSampler()
            else:  # map-style
                if shuffle:
                    sampler = RandomSampler(dataset, generator=generator)
                else:
                    sampler = SequentialSampler(dataset)

        if batch_size is not None and batch_sampler is None:
            # auto_collation without custom batch_sampler
            batch_sampler = BatchSampler(sampler, batch_size, drop_last)

        self.batch_size = batch_size
        self.drop_last = drop_last
        self.sampler = sampler
        self.batch_sampler = batch_sampler

        if collate_fn is None:
            if self._auto_collation:
                collate_fn = default_collate
            else:
                collate_fn = default_convert

        self.collate_fn = collate_fn
        self._IterableDataset_len_called = (
            None  # See NOTE [ IterableDataset and __len__ ] in ``__len__``
        )

        self._iterator = None

        # the sampler/data stream.  Mutating these while a worker iterator is
        # alive can otherwise silently mix epochs or strand queued indices.
        self.__initialized = True

        self.check_worker_number_rationality()

    def check_worker_number_rationality(self) -> None:
        # Warn when num_workers exceeds the logical CPUs available to this
        if not self.num_workers or self.num_workers == 0:
            return
        max_num_worker_suggest = _cpu_count()
        cpuset_checked = hasattr(os, "sched_getaffinity")
        if max_num_worker_suggest is None or self.num_workers > max_num_worker_suggest:
            warnings.warn(
                _worker_rationality_warning_msg(
                    max_num_worker_suggest, self.num_workers, cpuset_checked
                ),
                stacklevel=2,
            )

    @property
    def _auto_collation(self) -> bool:
        return self.batch_sampler is not None

    @property
    def _index_sampler(self) -> Sampler | Iterable:
        # The actual sampler used for generating indices to read data at each
        # time. This would be ``batch_sampler`` if in auto-collation mode, and
        # ``sampler`` otherwise. We can't change ``sampler`` and
        # ``batch_sampler`` attributes for BC reasons.
        if self._auto_collation:
            return self.batch_sampler
        else:
            return self.sampler

    @property
    def multiprocessing_context(self):
        return self.__multiprocessing_context

    @multiprocessing_context.setter
    def multiprocessing_context(self, multiprocessing_context):
        if multiprocessing_context is not None:
            if self.num_workers > 0:
                if isinstance(multiprocessing_context, str):
                    valid_start_methods = multiprocessing.get_all_start_methods()
                    if multiprocessing_context not in valid_start_methods:
                        raise ValueError(
                            "multiprocessing_context option "
                            f"should specify a valid start method in {valid_start_methods!r}, but got "
                            f"multiprocessing_context={multiprocessing_context!r}"
                        )
                    multiprocessing_context = multiprocessing.get_context(
                        multiprocessing_context
                    )

                if not isinstance(
                    multiprocessing_context, multiprocessing.context.BaseContext
                ):
                    raise TypeError(
                        "multiprocessing_context option should be a valid context "
                        "object or a string specifying the start method, but got "
                        f"multiprocessing_context={multiprocessing_context}"
                    )
            else:
                raise ValueError(
                    "multiprocessing_context can only be used with "
                    "multiprocessing (num_workers > 0)"
                )
        self.__multiprocessing_context = multiprocessing_context

    def __setattr__(self, attr, value):
        if self.__initialized and attr in (
            "batch_size",
            "batch_sampler",
            "sampler",
            "drop_last",
            "dataset",
            "persistent_workers",
        ):
            raise ValueError(
                f"{attr} attribute should not be set after "
                f"{self.__class__.__name__} is initialized"
            )
        super().__setattr__(attr, value)

    def _get_iterator(self):
        if self.num_workers == 0:
            return _SingleProcessDataLoaderIter(self)
        else:
            self.check_worker_number_rationality()
            return _MultiProcessDataLoaderIter(self)

    def __iter__(self) -> Iterator[Any]:
        if self.persistent_workers and self.num_workers > 0:
            if self._iterator is None:
                # Create a single copy of this iterator per `DataLoader`
                # instance, so it can be re-used across epochs.
                self._iterator = self._get_iterator()
            else:
                self._iterator._reset(self)
            return self._iterator
        else:
            return self._get_iterator()

    def __len__(self) -> int:
        if self._dataset_kind == _DatasetKind.Iterable:
            # NOTE [ IterableDataset and __len__ ]
            #
            # For `IterableDataset`, `__len__` could be inaccurate when one
            # naively does multi-processing data loading, since the samples
            # will be duplicated. However, no real use case should be actually
            # using that behavior, so it should count as a user error. We
            # should generally trust user code to do the proper thing (e.g.,
            # configure each replica differently in `__iter__`), and give us
            # the correct `__len__` if they choose to implement it (this will
            # still throw if the dataset does not implement a `__len__`).
            #
            # To provide a further warning, we track if `__len__` was called on
            # the `DataLoader`, save the returned value in
            # `_IterableDataset_len_called`, and warn if the iterator ends up
            # yielding more than this number of samples.

            # Cannot statically verify that dataset is Sized
            length = self._IterableDataset_len_called = len(self.dataset)
            if (
                self.batch_size is not None
            ):  # IterableDataset doesn't allow custom sampler or batch_sampler
                if self.drop_last:
                    length = length // self.batch_size
                else:
                    length = math.ceil(length / self.batch_size)
            return length
        else:
            return len(self._index_sampler)

    def _finalize_batch(self, data: Any) -> Any:
        # Post-processing applied in the main process after a batch has been
        # skipped when no accelerator is available.
        if self.pin_memory and tp.cuda.is_available():
            data = _pin_memory(data)
        return self._apply_device(data)

    def _apply_device(self, data: Any) -> Any:
        if self.device is not None:
            data = self._move_to_device(data)
        return data

    def _move_to_device(self, data: Any) -> Any:
        if isinstance(data, Tensor):
            if self.device == "cpu":
                return data.cpu()
            return data.cuda(self.device, non_blocking=self.pin_memory)
        if isinstance(data, tuple):
            return tuple(self._move_to_device(item) for item in data)
        if isinstance(data, list):
            return [self._move_to_device(item) for item in data]
        if isinstance(data, dict):
            return {key: self._move_to_device(value) for key, value in data.items()}
        return data
```