grain.experimental module

Contents

grain.experimental module#

Experimental Grain APIs.

List of Members#

grain.experimental.FlatMapTransform[source]#

alias of FlatMap

class grain.experimental.DatasetOptions(*, filter_warn_threshold_ratio=_Default(value=0.9), filter_raise_threshold_ratio=_Default(value=None), execution_tracking_mode=_Default(value=<ExecutionTrackingMode.DISABLED: 1>), min_shm_size=_Default(value=0))[source]#

Holds options used by dataset transformations.

This dataclass manages execution, telemetry, and performance parameters for PyGrain data pipelines. It tracks which fields are explicitly set by the user versus which fallback to default values, enabling intelligent option merging across different pipeline stages.

Parameters:

  • filter_warn_threshold_ratio (float | None | _Default[float])

  • filter_raise_threshold_ratio (float | None | _Default[None])

  • execution_tracking_mode (ExecutionTrackingMode | _Default[ExecutionTrackingMode])

  • min_shm_size (int | _Default[int])

filter_warn_threshold_ratio#

If the ratio of filtered out elements is above these thresholds, a warning will be issued. Value None disables the check. The ratio is calculated on non-overlapping windows of 1000 elements. For instance, with filter_warn_threshold_ratio=0.9 and 901 elements out of the first 1000 (or elements 1000…2000) filtered out, a warning will be issued.

Type:

float | None | grain._src.python.dataset.base._Default[float]

filter_raise_threshold_ratio#

If the ratio of filtered out elements is above these thresholds, an exception will be issued. Value None disables the check.

Type:

float | None | grain._src.python.dataset.base._Default[None]

execution_tracking_mode#

The collection of execution statistics like total processing time taken by each transformation, number of elements produced etc. can be managed through various modes. If DISABLED, no statistics are collected.If STAGE_TIMING, the time it takes to process each transormation is collected. See ExecutionTrackingMode for more details.

Type:

grain._src.python.dataset.base.ExecutionTrackingMode | grain._src.python.dataset.base._Default[grain._src.python.dataset.base.ExecutionTrackingMode]

min_shm_size#

The minimum size below which numpy arrays will copied between processes rather than passed via shared memory. For smaller arrays, the overhead of using shared memory can be higher than the cost of copying.

Type:

int | grain._src.python.dataset.base._Default[int]

Example

Applying custom options to dataset transformations:

import grain

ds = (
    grain.MapDataset.range(0, 1000)
    .filter(lambda x: x % 2 == 0)
    .to_iter_dataset()
)
# apply the DatasetOptions to create another IterDataset.
ds_options = grain.experimental.DatasetOptions(filter_raise_threshold_ratio=0.1)
ds = grain.experimental.WithOptionsIterDataset(ds, ds_options)
merge(other)[source]#

Merges these options with another DatasetOptions instance.

This merge logic respects explicit user configurations over defaults. Explicitly set options in self take highest precedence, followed by explicitly set options in other, followed by the class default values.

Parameters:

other (DatasetOptions | None) – Another DatasetOptions instance to merge into this one. If None, this method returns the current instance unmodified.

Returns:

A new DatasetOptions instance containing the merged

configuration values.

Return type:

DatasetOptions

Example

Demonstrating the explicit-set precedence during a merge:

import grain
# opt1 explicitly sets min_shm_size
opt1 = grain.experimental.DatasetOptions(min_shm_size=1024)

# opt2 explicitly sets min_shm_size AND filter thresholds
opt2 = grain.experimental.DatasetOptions(
    min_shm_size=512,
    filter_warn_threshold_ratio=None
)

# Merge opt2 into opt1. opt1's explicit values take precedence.
merged = opt1.merge(opt2)
assert merged.min_shm_size == 1024
assert merged.filter_warn_threshold_ratio is None
class grain.experimental.ExecutionTrackingMode(*values)[source]#

Represents different modes for tracking execution statistics.

Available modes:
DISABLED:

No execution statistics are measured. This mode is the default.

STAGE_TIMING:

The time taken for each transformation stage to execute is measured and recorded. This recorded time reflects the duration spent within the specific transformation to return an element, excluding the time spent in any parent transformations. The recorded time can be retrieved using grain.experimental.get_execution_summary method.

Example

To enable stage timing, set execution_tracking_mode in grain.experimental.DatasetOptions and pass it to your dataset pipeline using grain.experimental.WithOptionsIterDataset or grain.experimental.WithOptionsMapDataset:

import grain

options = grain.experimental.DatasetOptions(
    execution_tracking_mode=grain.experimental.ExecutionTrackingMode.STAGE_TIMING
)
ds = grain.MapDataset.range(10).to_iter_dataset()
ds_with_stage_timing = (
    grain.experimental.WithOptionsIterDataset(ds, options)
)

for element in ds_with_stage_timing:
  print(element)
grain.experimental.apply_transformations(ds, transformations)[source]#

Applies transformations to a dataset.

DEPRECATED: Use ds.apply(transformations) instead.

Parameters:
Returns:

Dataset of the same type with transformations applied.

Return type:

_ConsistentDatasetType

class grain.experimental.ElasticIterator(ds, global_batch_size, shard_options, *, read_options=ReadOptions(num_threads=16, prefetch_buffer_size=500), multiprocessing_options=None)[source]#

Iterator supporting recovery from a checkpoint after changes in sharding.

The input dataset is expected to be unbatched and unsharded. In order to provide elasticity guarantee this iterator includes both, batching and sharding. This iterator explicitly disallows many-to-one transformations without a fixed ratio, like filter and generic IterDataset transformations. The implementation differs for MapDatasets and IterDatasets.

MapDatasets:

The iterator supports elastic re-configuration by having each shard produce the same exact checkpoint (while producing different data) as long as they are advanced the same number of steps.

State of any shard can be used to restore the state of all of the shards after changes in sharding and global batch size.

IterDatasets:

IterDatasets support is still under development and comes with a few limitations. This class does not guarantee determinism between scaling. The limit of parallelism is the number of shards. The current implementation doesn’t support multiprocessing.

Example

Creating an elastic iterator for dynamic sharding recovery:

import grain

source_ds = grain.MapDataset.range(1000)
shard_opts = grain.sharding.ShardOptions(shard_index=0, shard_count=4)
elastic_iter = grain.experimental.ElasticIterator(
    ds=source_ds,
    global_batch_size=128,
    shard_options=shard_opts,
)

# The state from this iterator can now safely restore a job that
# restarts with, for example, 8 shards and a batch size of 256.
Parameters:
__init__(ds, global_batch_size, shard_options, *, read_options=ReadOptions(num_threads=16, prefetch_buffer_size=500), multiprocessing_options=None)[source]#

Initializes the ElasticIterator.

Parameters:
__next__()[source]#

For backwards compatibility with direct next() calls.

This allows using ElasticIterator as a one-shot iterator directly, delegating to a base iterator created on the first call.

Returns:

The next element in the iteration.

Return type:

Any

class grain.experimental.WithOptionsIterDataset(parent, options)[source]#

Applies options to transformations in the pipeline.

The options will apply to all transformations in the pipeline (before and after WithOptionsIterDataset). The options can be set multiple times in the pipeline, in which case they are merged. If the same option is set multiple times, the latest value takes precedence.

Example:

ds = MapDataset.range(5).to_iter_dataset()
ds = WithOptionsIterDataset(
      ds,
      DatasetOptions(
          filter_warn_threshold_ratio=0.6,
          filter_raise_threshold_ratio=0.8,
      ),
    )
ds = ds.filter(...)
ds = WithOptionsIterDataset(
      ds,
      DatasetOptions(filter_warn_threshold_ratio=0.7),
    )
ds = ds.filter(...)

In this case, the options will be:

filter_warn_threshold_ratio=0.7
filter_raise_threshold_ratio=0.8
Parameters:
__init__(parent, options)[source]#
Parameters:
__iter__()[source]#

Returns an iterator for the dataset.

Propagates the options through the pipeline ensuring there’s a single merged options object. Depending on the options, may apply iterator transformations.

Return type:

DatasetIterator[T]

class grain.experimental.ParquetIterDataset(path, **read_kwargs)[source]#

An IterDataset for a Parquet format file.

This dataset provides an iterator over records stored in Parquet files. It natively handles both single-file reads and multi-file interleaving. If a sequence of multiple paths is provided, the dataset automatically interleaves reads from the files (reading up to 16 files concurrently by default).

Additional keyword arguments provided during initialization are forwarded directly to the underlying pyarrow.parquet.ParquetFile constructor. This allows users to configure advanced Arrow features like memory mapping or custom buffer sizes.

Example

Initializing a dataset to read records from a Parquet file with memory_map option passed to ParquetFile:

import os
import tempfile
import grain
import pyarrow as pa
import pyarrow.parquet as pq

with tempfile.TemporaryDirectory() as tmpdir:
  tmp_path = os.path.join(tmpdir, "data.parquet")
  table = pa.table({"id": [1, 2], "val": ["A", "B"]})
  pq.write_table(table, tmp_path)

  # Create a Parquet dataset with a keyword arg.
  ds = grain.experimental.ParquetIterDataset(
      tmp_path,
      memory_map=True
  )

  # Print each record from the dataset.
  for record in ds:
    print(record)
Parameters:

path (str | Sequence[str])

__init__(path, **read_kwargs)[source]#

Initializes ParquetIterDataset.

Parameters:

  • path (str | Sequence[str]) – A path or sequence of paths to Parquet format files. If multiple paths are provided, they are interleaved with at most 16 files read concurrently.

  • **read_kwargs – Keyword arguments to pass to pyarrow.parquet.ParquetFile.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

class grain.experimental.TFRecordIterDataset(path)[source]#

An IterDataset for a TFRecord format file.

Iterates over a TFRecord file sequentially and yields records as raw bytes.

Example

Reading and iterating over elements from a TFRecord file:

import tempfile
import grain
import tensorflow as tf

with tempfile.TemporaryDirectory() as temp_dir:
  path = f"{temp_dir}/sample.tfrecord"

  with tf.io.TFRecordWriter(path) as writer:
    writer.write(b"record_0")
    writer.write(b"record_1")

  ds = grain.experimental.TFRecordIterDataset(path)
  ds_iter = iter(ds)

  print(next(ds_iter))
  # b'record_0'
  print(next(ds_iter))
  # b'record_1'
Parameters:

path (str)

__init__(path)[source]#

Initializes the TFRecordIterDataset.

Parameters:

path (str) – Path to the TFRecord file.

__iter__()[source]#

Returns an iterator over the TFRecord file.

Return type:

DatasetIterator[T]

grain.experimental.batch_and_pad(values, *, batch_size, pad_value=0)[source]#

Batches the given values and, if needed, pads the batch to the given size.

Can be passed to ds.batch as batch_fn to avoid the need to drop the remainder data and pad it instead.

Example usage:

ds = grain.MapDataset.range(1, 5)
batch_size = 3
batch_fn = functools.partial(
    grain.experimental.batch_and_pad, batch_size=batch_size)
ds = ds.batch(batch_size, batch_fn=batch_fn)
list(ds) == [np.ndarray([1, 2, 3]), np.ndarray([4, 0, 0])]
Parameters:

  • values (Sequence[T]) – The values to batch.

  • batch_size (int) – Target batch size. If the number of values is smaller than this, the batch is padded with pad_value to the given size.

  • pad_value (Any) – The value to use for padding.

Returns:

A batch of values with a new batch dimension at the front.

Return type:

T

class grain.experimental.CacheIterDataset(parent)[source]#

Caches elements of an IterDataset in memory.

Parameters:

parent (dataset.IterDataset[T])

__init__(parent)[source]#

Caches elements of an IterDataset in memory.

Parameters:

parent (IterDataset[T]) – The parent IterDataset whose elements are to be cached.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

class grain.experimental.FlatMapMapDataset(parent, transform)[source]#

Flat map for one-to-many split.

Wraps a parent MapDataset and applies a FlatMapTransform to each element. Supports random access by calculating a deterministic virtual length.

Note

If the FlatMapTransform returns fewer items than the specified max_fan_out, the dataset fills the remaining slots with None values. If the number of returned items exceeds max_fan_out, a ValueError is raised. Hence the max_fan_out should be greater than or equal to the maximum number of items returned by the FlatMapTransform for any element.

Example

Splitting strings into words using a custom FlatMapTransform:

import dataclasses
import grain

# Custom FlatMapTransform for text splitting.
@dataclasses.dataclass(frozen=True)
class SplitWords(grain.experimental.FlatMapTransform):
  max_fan_out: int

  def flat_map(self, sentence: str) -> list[str]:
    return sentence.split()

# Parent dataset with 2 elements
parent_ds = grain.MapDataset.source(["hello world", "grain"])

print(list(parent_ds))
# ['hello world', 'grain']

# Create a FlatMapMapDataset with max_fan_out=3
flatmap_map_ds = grain.experimental.FlatMapMapDataset(
    parent_ds, SplitWords(3)
)

# Virtual length is len(parent_ds) * max_fan_out = 2 * 3 = 6
assert len(flatmap_map_ds) == 6

# Elements are split and padded with None
# None values are not returned when converting to a list.
print(list(flatmap_map_ds))
# ['hello', 'world', 'grain']

# Elements can be accessed randomly by index:
for i in range(len(flatmap_map_ds)):
  print(f'{i=}: {flatmap_map_ds[i]}')
# i=0: hello
# i=1: world
# i=2: None
# i=3: grain
# i=4: None
# i=5: None
Parameters:
__getitem__(index)[source]#

Returns the element for the index or None if missing.

__init__(parent, transform)[source]#

Initializes the FlatMapMapDataset.

Parameters:

  • parent (MapDataset) – The parent MapDataset instance.

  • transform (FlatMap) – A FlatMapTransform instance that must implement the max_fan_out property.

class grain.experimental.FlatMapIterDataset(parent, transform)[source]#

Flat map for one-to-many split.

Wraps a parent IterDataset and applies a FlatMapTransform to each element, yielding the resulting sequence as a stream.

Example

Splitting strings into words using a custom FlatMapTransform:

import grain

# Create a custom FlatMapTransform
class DuplicateString(grain.experimental.FlatMapTransform):

  def flat_map(self, element: str) -> list[str]:
    return [element, element.upper()]

# Create a parent IterDataset
parent_ds = grain.MapDataset.source(
    ["hello", "grain"]
).to_iter_dataset()

# Apply the FlatMapTransform to the IterDataset
flatmap_iter_ds = grain.experimental.FlatMapIterDataset(
    parent_ds, DuplicateString()
)

# Iterating through an IterDataset yields only the actual elements.
iterator = iter(flatmap_iter_ds)
for _ in range(len(list(flatmap_iter_ds))):
  print(next(iterator))
# hello
# HELLO
# grain
# GRAIN
Parameters:
__init__(parent, transform)[source]#

Initializes the FlatMapIterDataset.

Parameters:

  • parent (IterDataset) – The parent IterDataset instance.

  • transform (FlatMap) – A FlatMapTransform instance that must implement the flat_map method.

__iter__()[source]#

Returns an iterator for this dataset.

class grain.experimental.InterleaveIterDataset(datasets, *, cycle_length, num_make_iter_threads=1, make_iter_buffer_size=1, iter_buffer_size=1)[source]#

Interleaves the given sequence of datasets.

The sequence can be a MapDataset.

Concurrently processes at most cycle_length iterators and interleaves their elements. If cycle_length is larger than the number of datasets, then the behavior is similar to mixing the datasets with equal proportions. If cycle_length is 1, the datasets are chained.

This dataset can be combined with mp_prefetch to parallelize reads from sources that do not support random access.

Element spec inference assumes that all input datasets have the same element spec.

Example

Interleaving four datasets with two active iterators:

import grain

def make_source(start):
  return grain.MapDataset.range(start, start + 2).to_iter_dataset()

sources = grain.MapDataset.source([0, 10, 20, 30]).map(make_source)

print(list(sources[1]))
# [10, 11]

interleaved_ds = grain.experimental.InterleaveIterDataset(
    sources,
    cycle_length=2,
)

print(list(interleaved_ds))
# [0, 10, 1, 11, 20, 30, 21, 31]
Parameters:

  • datasets (Sequence[IterDataset[T] | MapDataset[T]])

  • cycle_length (int | AutotuneParameter)

  • num_make_iter_threads (int)

  • make_iter_buffer_size (int)

  • iter_buffer_size (int)

__init__(datasets, *, cycle_length, num_make_iter_threads=1, make_iter_buffer_size=1, iter_buffer_size=1)[source]#

Initializes the InterleaveIterDataset.

Parameters:

  • datasets (Sequence[IterDataset[T] | MapDataset[T]]) – A sequence of IterDataset or MapDataset objects, or a MapDataset of datasets to be interleaved.

  • cycle_length (int | AutotuneParameter) – The maximum number of input datasets from which elements will be processed concurrently. If cycle_length is greater than the total number of datasets, all available datasets will be interleaved. If cycle_length is 1, the datasets will be processed sequentially.

  • num_make_iter_threads (int) – Optional. The number of threads to use for asynchronously creating new iterators and starting prefetching elements (for each iterator) from the underlying datasets. Default value is 1, with this we’ll create one background thread to asynchronously create iterators.

  • make_iter_buffer_size (int) – Optional. The number of iterators to create and keep ready in advance in each preparation thread. This helps in reducing latency by ensuring iterators are available when needed. Default value is 1, with this we’ll always keep the next iterator ready in advance.

  • iter_buffer_size (int) – Optional. The number of elements to prefetch from each iterator. Default value is 1.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

class grain.experimental.LimitIterDataset(parent, count)[source]#

Limits the number of elements in the dataset.

Iteration stops after count elements have been produced or when the parent dataset is exhausted, whichever occurs first.

Example

Limiting a dataset to two elements:

import grain

parent_ds = grain.MapDataset.range(5).to_iter_dataset()

print(list(parent_ds))
# [0, 1, 2, 3, 4]

limited_ds = grain.experimental.LimitIterDataset(
    parent_ds,
    count=2,
)

print(list(limited_ds))
# [0, 1]
Parameters:
__init__(parent, count)[source]#

Initializes the LimitIterDataset.

Parameters:

  • parent (IterDataset[T]) – The dataset to limit.

  • count (int) – The maximum number of elements to include in the dataset.

Raises:

ValueError – If count is not positive.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

_LimitDatasetIterator[T]

class grain.experimental.RngPool(seed)[source]#

RNG pool.

Parameters:

seed (int)

acquire_rng(index, *, op_seed=0)[source]#

Acquire RNG.

Parameters:

  • index (int)

  • op_seed (int)

Return type:

numpy.random.Generator

class grain.experimental.FirstFitPackIterDataset(parent, *, length_struct, num_packing_bins, seed=0, shuffle_bins=True, shuffle_bins_group_by_feature=None, meta_features=(), pack_alignment_struct=None, padding_struct=None, max_sequences_per_bin=None)[source]#

Implements first-fit packing of sequences.

Packing, compared to concat-and-split, avoids splitting sequences by padding instead. Larger number of packing bins reduce the amount of padding. If the number of bins is large, this can cause epoch leakage (data points from multiple epochs getting packed together).

This uses a simple first-fit packing algorithm that: 1. Creates N bins. 2. Adds elements (in the order coming from the parent) to the first bin that has enough space. 3. Once an element doesn’t fit, emits all N bins as elements. 4. (optional) Shuffles bins. 5. Loops back to 1 and starts with the element that didn’t fit.

Example

Pack variable-length sequences into fixed-length outputs:

import grain
import numpy as np

# Parent dataset with variable length sequences.
parent_ds = grain.MapDataset.source([
    {"x": np.array([1, 2])},
    {"x": np.array([3, 4, 5])},
    {"x": np.array([6])},
    {"x": np.array([7, 8])},
]).to_iter_dataset()

# The first element of the parent dataset has "x" with shape (2,).
parent_ds_iterator = iter(parent_ds)
parent_ds_first_element = next(parent_ds_iterator)
print(parent_ds_first_element["x"])
# [1 2]

# Pack sequences into 2 bins with target length 4 for feature "x".
packed_ds = grain.experimental.FirstFitPackIterDataset(
  parent_ds,
  length_struct={"x": 4},
  num_packing_bins=2,
  shuffle_bins=False,
)

# Now the first and the third elements of parent dataset are packed into
# one bin of shape (4,)
packed_ds_iterator = iter(packed_ds)
packed_ds_first_element = next(packed_ds_iterator)
print(packed_ds_first_element["x"].shape)
# (4,)
print(packed_ds_first_element["x"])
# [1 2 6 0]
Parameters:

  • parent (IterDataset)

  • length_struct (Any)

  • num_packing_bins (int)

  • seed (int)

  • shuffle_bins (bool)

  • shuffle_bins_group_by_feature (str | None)

  • meta_features (Sequence[str])

  • pack_alignment_struct (Any)

  • padding_struct (Any)

  • max_sequences_per_bin (int | None)

__init__(parent, *, length_struct, num_packing_bins, seed=0, shuffle_bins=True, shuffle_bins_group_by_feature=None, meta_features=(), pack_alignment_struct=None, padding_struct=None, max_sequences_per_bin=None)[source]#

Creates a dataset that packs sequences using the first-fit strategy.

Parameters:

  • parent (IterDataset) – Parent dataset with variable length sequences.

  • length_struct (Any) – Target sequence length for each feature.

  • num_packing_bins (int) – Number of bins to pack sequences into.

  • seed (int) – Random seed for shuffling bins.

  • shuffle_bins (bool) – Whether to shuffle bins after packing.

  • shuffle_bins_group_by_feature (str | None) – Feature to group by for shuffling.

  • meta_features (Sequence[str]) – Meta features that do not need packing logic. They can be sequence meta-features (if present in length_struct, packed and padded to target length) or non-sequence meta-features (if not present in length_struct, returned as a list of meta-features from the packed examples in each bin).

  • pack_alignment_struct (Any) – Optional per-feature alignment values.

  • padding_struct (Any) – Optional per-feature padding values.

  • max_sequences_per_bin (int | None) – Optional maximum number of input sequences that can be packed into a bin

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator

class grain.experimental.BestFitPackIterDataset(parent, *, length_struct, num_packing_bins, seed=0, shuffle_bins=True, shuffle_bins_group_by_feature=None, meta_features=(), pack_alignment_struct=None, padding_struct=None, max_sequences_per_bin=None)[source]#

Implements best-fit packing of sequences.

The best-fit algorithm attempts to pack elements more efficiently than first-fit by placing each new element into the bin that will leave the smallest remaining space (i.e., the “tightest” fit). This can lead to less overall padding compared to the simpler first-fit approach, especially when element sizes vary significantly.

Example

Pack variable-length sequences into fixed-length outputs:

import grain
import numpy as np

# Parent dataset with variable length sequences.
parent_ds = grain.MapDataset.source([
    {"x": np.array([1, 2])},
    {"x": np.array([3, 4, 5])},
    {"x": np.array([6])},
    {"x": np.array([7, 8])},
]).to_iter_dataset()

# The first element of the parent dataset has "x" with shape (2,).
parent_ds_iterator = iter(parent_ds)
parent_ds_first_element = next(parent_ds_iterator)
print(parent_ds_first_element["x"])
# [1, 2]

# Best-fit packing with 2 bins and target length 4.
packed_ds = grain.experimental.BestFitPackIterDataset(
    parent_ds,
    length_struct={"x": 4},
    num_packing_bins=2,
    shuffle_bins=False,
)

# Now the first and the fourth elements of parent dataset are packed into
# one bin of shape (4,)
packed_ds_iterator = iter(packed_ds)
packed_ds_first_element = next(packed_ds_iterator)
print(packed_ds_first_element["x"].shape)
# (4,)
print(packed_ds_first_element["x"])
# [1 2 7 8]
Parameters:

  • parent (IterDataset)

  • length_struct (Any)

  • num_packing_bins (int)

  • seed (int)

  • shuffle_bins (bool)

  • shuffle_bins_group_by_feature (str | None)

  • meta_features (Sequence[str])

  • pack_alignment_struct (Any)

  • padding_struct (Any)

  • max_sequences_per_bin (int | None)

__init__(parent, *, length_struct, num_packing_bins, seed=0, shuffle_bins=True, shuffle_bins_group_by_feature=None, meta_features=(), pack_alignment_struct=None, padding_struct=None, max_sequences_per_bin=None)[source]#

Creates a dataset that packs sequences using the best-fit strategy.

Parameters:

  • parent (IterDataset) – Parent dataset with variable length sequences.

  • length_struct (Any) – Target sequence length for each feature.

  • num_packing_bins (int) – Number of bins to pack sequences into.

  • seed (int) – Random seed for shuffling bins.

  • shuffle_bins (bool) – Whether to shuffle bins after packing.

  • shuffle_bins_group_by_feature (str | None) – Feature to group by for shuffling.

  • meta_features (Sequence[str]) – Meta features that do not need packing logic. They can be sequence meta-features (if present in length_struct, packed and padded to target length) or non-sequence meta-features (if not present in length_struct, returned as a list of meta-features from the packed examples in each bin).

  • pack_alignment_struct (Any) – Optional per-feature alignment values.

  • padding_struct (Any) – Optional per-feature padding values.

  • max_sequences_per_bin (int | None) – Optional maximum number of input sequences that can be packed into a bin

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator

class grain.experimental.BOSHandling(*values)[source]#

The BOS handling done inside a packing algorithm.

class grain.experimental.ConcatThenSplitIterDataset(parent, *, length_struct, meta_features=(), split_full_length_features=True, bos_handling=BOSHandling.DO_NOTHING, bos_features=(), bos_token_id=None)[source]#

Implements concat-then-split packing for sequence features.

This assumes that elements of the parent dataset are unnested dictionaries and entries are either scalars or NumPy arrays. The first dimension is considered the sequence dimension and its size may vary between elements. All other dimensions must be the same size for all elements. Scalars are treated as 1-dimensional arrays of size 1.

On a high level this concatenates the underlying dataset and then splits it at target sequence lengths intervals. This is well defined for the case of a single feature. For multiple features we start with an empty buffer and concatenate elements until at least one feature is fully packed. As an optimization, elements from the parent dataset that are already fully packed are passed through in priority. When the buffer contains enough elements to fill at least one feature to its target sequence length, we pack the buffer. The last element might not fully fit and will be split. The remainder of the split stays in the buffer.

When packing features we also create {feature_name}_positions and {feature_name}_segment_ids features. They are 1D arrays of size sequence_length. Segment IDs start at 1 and enumerate the elements of the packed element. Positions indicate the position within the unpacked sequence.

Features can be “meta features” in which case they are never split and we do not create *_positions and *_segment_ids features for them.

Example

Applying concat-then-split packing to a dataset with variable-length sequence features:

import grain
import numpy as np

# Input data with varying sequence lengths
data = [
    {"inputs": np.array([1, 2, 3]), "targets": np.array([10, 20])},
    {"inputs": np.array([4, 5, 6]), "targets": np.array([30, 40])},
]
parent_ds = grain.MapDataset.source(data).to_iter_dataset()

parent_ds_iterator = iter(parent_ds)
parent_ds_first_batch = next(parent_ds_iterator)

print(parent_ds_first_batch["inputs"])
# array([1, 2, 3])

# Pack features into specific lengths. The dataset buffers elements
# until a feature reaches its target length.
packed_ds = grain.experimental.ConcatThenSplitIterDataset(
    parent=parent_ds,
    length_struct={"inputs": 4, "targets": 3},
)

packed_ds_iter = iter(packed_ds)
packed_ds_first_batch = next(packed_ds_iter)

print(packed_ds_first_batch["inputs"])
# array([1, 2, 3, 4])

print(packed_ds_first_batch["inputs_segment_ids"])
# array([1, 1, 1, 2])

print(packed_ds_first_batch["inputs_positions"])
# array([0, 1, 2, 0])
Parameters:

  • parent (dataset.IterDataset)

  • length_struct (Mapping[str, int])

  • meta_features (Collection[str])

  • split_full_length_features (bool)

  • bos_handling (BOSHandling)

  • bos_features (Collection[str])

  • bos_token_id (int | None)

__init__(parent, *, length_struct, meta_features=(), split_full_length_features=True, bos_handling=BOSHandling.DO_NOTHING, bos_features=(), bos_token_id=None)[source]#

Creates a dataset that concat-then-splits sequences from the parent.

Parameters:

  • parent (IterDataset) – The parent dataset.

  • length_struct (Mapping[str, int]) – Mapping from feature name to target sequence length.

  • meta_features (Collection[str]) – Set of feature names that are considered meta features. Meta features are never split and will be duplicated when other features of the same element are split. Otherwise, meta features are packed normally (they have their own sequence length). No *_positions and *_segment_ids features are created for meta features.

  • split_full_length_features (bool) – Whether full-length features are split, or they are considered packed and passed through in priority. Setting split_full_length_features=False is an optimization when some sequences already have the target length, and you don’t want them to be split. This optimization is not used by default.

  • bos_handling (BOSHandling) – The instructions for handling BOS tokens (by default, no BOS token is added).

  • bos_features (Collection[str]) – The features to which BOS handling is applied in case BOS is used.

  • bos_token_id (int | None) – The token indicating BOS in case BOS is used.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator

grain.experimental.multithread_prefetch(ds, num_threads, buffer_size, sequential_slice=False)[source]#

Uses a pool of threads to prefetch elements ahead of time.

This is a thread-based alternative to multiprocess_prefetch intended to be used with free-threaded Python.

It works by sharding the input dataset into num_threads shards, and interleaving them. Each shard is read by a separate thread inside InterleaveIterDataset.

Parameters:

  • ds (IterDataset[T]) – The parent dataset to prefetch from.

  • num_threads (int) – The number of threads to use for prefetching. If 0, prefetching is disabled and this is a no-op.

  • buffer_size (int) – The size of the prefetch buffer for each thread.

  • sequential_slice (bool) – Whether to use sequential slicing.

Returns:

An IterDataset that prefetches elements from ds using multiple threads.

Return type:

IterDataset[T]

class grain.experimental.ThreadPrefetchIterDataset(parent, *, prefetch_buffer_size)[source]#

Iterable dataset that uses a synchronized queue for prefetching.

This is a thread-based alternative to MultiprocessPrefetchIterDataset.

Example

Prefetching elements from an iterable dataset:

import grain

parent_ds = grain.MapDataset.range(4).to_iter_dataset()
print(list(parent_ds))
# [0, 1, 2, 3]

prefetched_ds = grain.experimental.ThreadPrefetchIterDataset(
    parent_ds,
    prefetch_buffer_size=2,
)

print(list(prefetched_ds))
# [0, 1, 2, 3]
Parameters:
__init__(parent, *, prefetch_buffer_size)[source]#

Initializes a ThreadPrefetchIterDataset.

Parameters:

  • parent (IterDataset[T]) – The parent dataset to prefetch from.

  • prefetch_buffer_size (int | AutotuneParameter) – The size of the prefetch buffer. Must be greater than or equal to 0. If 0, prefetching is disabled and this is a noop.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

class grain.experimental.ThreadPrefetchDatasetIterator(parent, prefetch_buffer_size)[source]#

Iterator that performs prefetching using a synchronized queue.

This iterator wraps a checkpointable iterator and asynchronously fetches future elements using a background thread. Prefetched elements are stored in an in-memory queue and returned when requested by the consumer.

The iterator preserves the parent iterator state and supports checkpointing through get_state() and set_state().

Example

Demonstrating state restoration with a prefetch buffer:

import grain

parent_ds = grain.MapDataset.range(4).to_iter_dataset()
parent_iter = iter(parent_ds)

ds_iter = grain.experimental.ThreadPrefetchDatasetIterator(
    parent_iter,
    prefetch_buffer_size=2,
)

print(next(ds_iter))
# 0

state = ds_iter.get_state()

print(next(ds_iter))
# 1

print(next(ds_iter))
# 2

ds_iter.set_state(state)

print(next(ds_iter))
# 1
Parameters:

  • parent (CheckpointableIterator[T])

  • prefetch_buffer_size (int | grain_options.AutotuneParameter)

__init__(parent, prefetch_buffer_size)[source]#

Initializes a ThreadPrefetchDatasetIterator.

Parameters:

  • parent (CheckpointableIterator[T]) – The checkpointable iterator to prefetch from.

  • prefetch_buffer_size (int | AutotuneParameter) – The size of the prefetch buffer. If 0, prefetching is disabled.

__next__()[source]#

Return the next item from the iterator. When exhausted, raise StopIteration

class grain.experimental.RebatchIterDataset(parent, batch_size, drop_remainder=False)[source]#

Rebatches the input PyTree elements.

This transformation takes PyTree elements with a leading batch dimension, combines data across consecutive input batches as needed, and slices them to produce output batches of the requested batch_size.

If drop_remainder is False (default), the final batch will contain any remaining elements (fewer than batch_size). If drop_remainder is True, the last batch is dropped if it contains fewer than batch_size elements.

Example

A common pattern for changing the batch size of an existing iterable dataset:

import grain
import numpy as np

# Create a dataset where each element is already a batch of size 2.
parent_ds = grain.MapDataset.range(1, 9).batch(2).to_iter_dataset()

print(list(parent_ds))
# [array([1, 2]), array([3, 4]), array([5, 6]), array([7, 8])]

# Rebatch the input batches to a new batch size.
rebatched_ds = grain.experimental.RebatchIterDataset(
    parent_ds,
    batch_size=3,
    drop_remainder=False,
)

iterator = iter(rebatched_ds)

# The first output batch combines rows from multiple input batches to
# produce the requested batch size.
batch1 = next(iterator)

print(batch1)
# array([1, 2, 3])

batch2 = next(iterator)

print(batch2)
# array([4, 5, 6])

batch3 = next(iterator)

print(batch3)
# array([7, 8])
Parameters:
__init__(parent, batch_size, drop_remainder=False)[source]#

An IterDataset that rebatches elements.

Parameters:

  • parent (IterDataset) – The parent IterDataset whose elements are to be rebatched.

  • batch_size (int) – The number of elements to batch together.

  • drop_remainder (bool) – Whether to drop the last batch if it is smaller than batch_size.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator

class grain.experimental.RepeatIterDataset(parent, num_epochs=None)[source]#

Repeats the underlying dataset for num_epochs.

If num_epochs is None, repeats indefinitely. Note that unlike RepeatMapDataset, RepeatIterDataset does not support re-seeding for each epoch. Each epoch will be identical.

Example

Repeat an iterable dataset for a fixed number of epochs:

import grain

# Build a small parent pipeline with 3 elements.
parent_ds = grain.MapDataset.range(9).batch(3).to_iter_dataset()

print(list(parent_ds))
# [array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8])]

# Repeat the dataset for exactly 2 epochs.
repeated_ds = grain.experimental.RepeatIterDataset(
    parent_ds,
    num_epochs=2,
)
print(list(repeated_ds))
# [array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]),
#  array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8])]
Parameters:
__init__(parent, num_epochs=None)[source]#

Initializes the RepeatIterDataset.

Parameters:

  • parent (IterDataset[T]) – The parent dataset.

  • num_epochs (int | None) – The number of times to repeat the dataset. If None, repeats indefinitely.

Raises:

ValueError – If num_epochs is not positive.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

_RepeatDatasetIterator[T]

class grain.experimental.WindowShuffleMapDataset(parent, *, window_size, seed)[source]#

Shuffles the parent dataset within a given window.

Shuffles the retrieval index within a range, given by window_size. Each unique index corresponds to exactly one shuffled index (i.e. there is a one-to-one mapping and hence a guarantee that no shuffled indices are repeated within a given window).

Example

Applying deterministic window-based shuffling to a dataset:

import grain

# Create a source dataset with consecutive elements.
parent_ds = grain.MapDataset.range(12)

print(list(parent_ds))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

# Shuffle elements independently within windows of size 3.
shuffled_ds = grain.experimental.WindowShuffleMapDataset(
    parent_ds,
    window_size=3,
    seed=42,
)

# Shuffled dataset
print(list(shuffled_ds))
# [0, 1, 2, 5, 3, 4, 6, 7, 8, 11, 9, 10]
Parameters:
__getitem__(index)[source]#

Returns the element for the index or None if missing.

__init__(parent, *, window_size, seed)[source]#

Initializes the WindowShuffleMapDataset.

Parameters:

  • parent (MapDataset) – The parent MapDataset to shuffle.

  • window_size (int) – The number of consecutive elements in each shuffle window.

  • seed (int) – Seed used to deterministically shuffle the elements within each window.

class grain.experimental.WindowShuffleIterDataset(parent, *, window_size, seed)[source]#

Shuffles the parent dataset within a given window.

Fetches window_size elements from the parent iterator and returns them in shuffled order. Each window is shuffled with different seed derived from the input seed.

Example

Applying window-based shuffling to an iterable dataset:

import grain

# Create a source dataset with consecutive elements.
parent_ds = grain.MapDataset.range(12).to_iter_dataset()

print(list(parent_ds))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

# Shuffle elements independently within windows of size 3.
shuffled_ds = grain.experimental.WindowShuffleIterDataset(
    parent_ds,
    window_size=3,
    seed=42,
)

# Shuffled dataset.
print(list(shuffled_ds))
# [2, 1, 0, 4, 3, 5, 8, 7, 6, 10, 9, 11]
Parameters:
__init__(parent, *, window_size, seed)[source]#

Initializes the WindowShuffleIterDataset.

Parameters:

  • parent (IterDataset) – The parent IterDataset to shuffle.

  • window_size (int) – The number of consecutive elements to read and shuffle in each window.

  • seed (int) – Seed used to deterministically shuffle the elements within each window.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

class grain.experimental.ZipMapDataset(parents)[source]#

Combines MapDatasets of the same length to return a tuple of items.

At each index, returns a tuple containing the corresponding element from each parent dataset. All parent datasets must have the same length; otherwise, raises a ValueError.

Examples

Combining corresponding elements from multiple map-style datasets:

import grain

# Create two source datasets of equal length.
inputs_ds = grain.MapDataset.source([10, 20, 30])
labels_ds = grain.MapDataset.source([40, 50, 60])

# Combine corresponding elements from both datasets.
zipped_ds = grain.experimental.ZipMapDataset([inputs_ds, labels_ds])

print(zipped_ds[0])
# (10, 40)

print(zipped_ds[1])
# (20, 50)
Parameters:

parents (Sequence[dataset.MapDataset[T]])

__getitem__(index)[source]#

Returns the element for the index or None if missing.

__init__(parents)[source]#

Initializes the ZipMapDataset.

Parameters:

parents (Sequence[MapDataset[T]]) – A sequence of MapDatasets to combine. All parent datasets must have the same length.

Raises:

  • ValueError – If no parent datasets are provided.

  • ValueError – If the parent datasets do not all have the same length.

class grain.experimental.ZipIterDataset(parents, *, strict=True)[source]#

Combines IterDatasets of the same length to return a tuple of items.

At each iteration, returns a tuple containing the next element from each parent dataset. By default (strict=True), all parent iterators are expected to produce the same number of elements; otherwise, a ValueError is raised during iteration. When strict=False, iteration stops when the shortest parent iterator is exhausted, matching the behavior of Python’s built-in zip.

Example

Iterating over corresponding elements from multiple datasets:

import grain

# Create two parent pipelines of equal length.
inputs_ds = grain.MapDataset.source([10, 20, 30]).to_iter_dataset()
labels_ds = grain.MapDataset.source([40, 50, 60]).to_iter_dataset()

# Combine corresponding elements from both pipelines.
zipped_ds = grain.experimental.ZipIterDataset([inputs_ds, labels_ds])
iterator = iter(zipped_ds)

print(next(iterator))
# (10, 40)

print(next(iterator))
# (20, 50)
Parameters:
__init__(parents, *, strict=True)[source]#

Initializes the ZipIterDataset.

Parameters:

  • parents (Sequence[IterDataset[T]]) – A sequence of IterDatasets to combine.

  • strict (bool) – If True (default), raises a ValueError during iteration when the parent iterators do not produce the same number of elements. If False, iteration stops when the shortest parent iterator is exhausted.

Raises:

ValueError – If no parent dataset is provided.

__iter__()[source]#

Returns an iterator for this dataset.

Return type:

DatasetIterator[T]

grain.experimental.index_shuffle()#
grain.experimental.assert_equal_output_after_checkpoint(ds)[source]#

Tests restoring an iterator to various checkpointed states.

Parameters:

ds (Any) – The dataset to test. It is recommended to use a small dataset, potentially created using grain.python.experimental.LimitIterDataset, to restrict the number of steps being tested. The underlying dataset iterator must implement get_state and set_state for checkpointing.

grain.experimental.device_put(ds, device, *, cpu_buffer_size=4, device_buffer_size=2)[source]#

Moves the data to the given devices with prefetching.

Stage 1: A CPU-side prefetch buffer. Stage 2: Per-device buffers for elements already transferred to the device.

Parameters:

  • ds (IterDataset) – Dataset to prefetch.

  • device – same arguments as in jax.device_put.

  • cpu_buffer_size (int) – Number of elements to prefetch on CPU.

  • device_buffer_size (int) – Number of elements to prefetch per device.

Returns:

Dataset with the elements prefetched to the devices.

Return type:

IterDataset

class grain.experimental.PerformanceConfig(multiprocessing_options: grain._src.python.options.MultiprocessingOptions | None = None, read_options: grain._src.python.options.ReadOptions | None = None)[source]#
Parameters:
grain.experimental.pick_performance_config(ds, *, ram_budget_mb, max_workers, max_buffer_size, samples_to_check=5)[source]#

Analyzes element size to choose an optimal number of workers for a MultiprocessPrefetchIterDataset.

Parameters:

  • ds (IterDataset) – The input dataset.

  • ram_budget_mb (int | None) – The user predicted RAM budget in megabytes.

  • max_workers (int | None) – The maximum number of processes to use.

  • max_buffer_size (int | None) – The maximum buffer size to use.

  • samples_to_check (int) – The number of samples to check to estimate element size.

Returns:

A PerformanceConfig object containing the optimal number of workers.

Return type:

PerformanceConfig

grain.experimental.get_element_spec(ds)[source]#

Returns specification of the elements produced by this dataset.

Does not instantiate iterator, perform any data reads or transformations.

Parameters:

ds (MapDataset | IterDataset) – MapDataset or IterDataset to get the element spec from.

Return type:

Any

grain.experimental.set_next_index(ds_iter, index)[source]#

Sets the next index for the dataset iterator.

Parameters:
Return type:

None

grain.experimental.get_next_index(ds_iter)[source]#

Returns the next index for the dataset iterator.

Parameters:

ds_iter (DatasetIterator)

Return type:

int

Contents