import grain.python as grain
import numpy as np

# DataSource
class MySource(grain.RandomAccessDataSource):
    def __len__(self):
        return len(self._data)
    def __getitem__(self, index):
        return self._data[index]

# Sampler
sampler = grain.IndexSampler(
    num_records=len(source),
    shuffle=True,
    num_epochs=None,
    seed=42,
    shard_options=grain.ShardByJaxProcess(drop_remainder=True),
)

# Deterministic transform
class MyMap(grain.MapTransform):
    def map(self, element):
        return element

# Random transform
class MyRandomMap(grain.RandomMapTransform):
    def random_map(self, element, rng):
        return element

# Assembling the pipeline
data_loader = grain.DataLoader(
    data_source=source,
    sampler=sampler,
    operations=[MyMap(), MyRandomMap(), grain.Batch(64, drop_remainder=True)],
    worker_count=4,
)

data_iterator = iter(data_loader)
batch = next(data_iterator)

        __
        __