grain.transforms module

Data transformation APIs.

List of Members

class grain.transforms.Batch(batch_size, drop_remainder=False, batch_fn=None)

Transformation that groups consecutive elements into batches.

Refer to grain.MapDataset.batch() or grain.IterDataset.batch() documentation for more details.

Parameters:

• batch_size (int)

• drop_remainder (bool)

• batch_fn (Callable[[Sequence[Any]], Any] | None)

class grain.transforms.Filter

Abstract base class for filter transformations for individual elements.

The pipeline will drop any element for which the filter function returns False.

Implementations should be threadsafe since they are often executed in parallel.

Example

Filtering out odd numbers from a dataset:

import grain

# Define a filter class that keeps only even numbers.
class KeepEven(grain.transforms.Filter):
  def filter(self, element: int) -> bool:
    return element % 2 == 0

# Create a parent dataset.
parent_ds = grain.MapDataset.range(10)
print(list(parent_ds))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

# Apply the filter transform.
transformed_ds = parent_ds.filter(KeepEven())
print(list(transformed_ds))
# [0, 2, 4, 6, 8]

class grain.transforms.Map

Abstract base class for all 1:1 transformations of elements.

Implementations should be threadsafe since they are often executed in parallel.

Example

Applying a custom map transformation to increment each element:

import grain

# Define a custom map transformation to add one to each element.
class AddOne(grain.transforms.Map):
  def map(self, element: int) -> int:
    return element + 1

# Create a parent dataset.
parent_ds = grain.MapDataset.range(5)
print(list(parent_ds))
# [0, 1, 2, 3, 4]

# Apply the map transformation.
transformed_ds = parent_ds.map(AddOne())
print(list(transformed_ds))
# [1, 2, 3, 4, 5]

class grain.transforms.MapWithIndex

Abstract base class for 1:1 transformations of elements and their index.

Implementations should be threadsafe since they are often executed in parallel.

Example

Applying a transformation that pairs each element with its index:

import grain

# Define a custom transform that inherits from MapWithIndex and implements
# the map_with_index method.
class AddIndex(grain.transforms.MapWithIndex):
  def map_with_index(self, index: int, element: int) -> tuple[int, int]:
    return index + element * 10

# Create a parent dataset.
parent_ds = grain.MapDataset.range(3)
print(list(parent_ds))
# [0, 1, 2]

# Apply the MapWithIndex transformation to the parent dataset.
transformed_ds = parent_ds.map_with_index(AddIndex())
print(list(transformed_ds))
# [0, 11, 22]

class grain.transforms.RandomMap

Abstract base class for all random 1:1 transformations of elements.

Implementations should be threadsafe since they are often executed in parallel.

Example

Applying a random map transformation to add noise:

import grain
import numpy as np

# Define a custom random map transformation to add random offset.
class AddRandomOffset(grain.transforms.RandomMap):
  def random_map(self, element: int, rng: np.random.Generator) -> int:
    return element + int(rng.integers(0, 10))

# Create a parent dataset and seed it for deterministic results.
parent_ds = grain.MapDataset.range(5).seed(42)
print(list(parent_ds))
# [0, 1, 2, 3, 4]

# Apply the random map transformation.
transformed_ds = parent_ds.random_map(AddRandomOffset())
print(list(transformed_ds))
# [2, 5, 5, 4, 9]

grain.transforms.Transformation

alias of Batch | Map | RandomMap | TfRandomMap | Filter | FlatMap | MapWithIndex

grain.transforms.Transformations

alias of Sequence[Batch | Map | RandomMap | TfRandomMap | Filter | FlatMap | MapWithIndex]

class grain.transforms.DatasetSelectionMap

Map from index to (constituent dataset index, index within dataset).

This abstract base class defines the interface for mapping a single global sequence index across multiple underlying (constituent) datasets. It acts as a routing table for mixed or concatenated data pipelines.

Note, this must be stateless, picklable and should avoid randomness to support determinism since it may be created and called concurrently in multiple processes.

Example

Implementing a simple concatenation map for two datasets of size 2 and 3:

import grain

class ConcatMap(grain.transforms.DatasetSelectionMap):
  def __len__(self) -> int:
    return 7
  def __getitem__(self, index: int) -> tuple[int, int]:
    if index >= len(self):
      raise IndexError("Index out of range")
    if index < 3:
      return (0, index)
    else:
      return (1, index - 3)

cmap = ConcatMap()
assert len(cmap) == 7
assert cmap[3] == (1, 0)