[tab] utils/dataloader.py (python)
import jax
import numpy as np
import grain
from typing import Any


# Note on Data Format:
# Grain's preferred file format is ArrayRecord. This implementation assumes you
# have converted your TFRecord files to ArrayRecord files. A simple script
# using `tf.data` to read TFRecords and `grain.ArrayRecordWriter` to write
# them can perform this one-time conversion. Each record in the ArrayRecord file
# is expected to be a dictionary with keys like "raw_video" and "sequence_length",
# mirroring the structure of the original TFRecord features.


class ProcessEpisodeAndSlice(grain.transforms.RandomMap):
    """
    A Grain Transformation that combines parsing, slicing, and normalizing.

    This class replaces the functionality of both `_parse_tfrecord_fn` and
    `_tf_process_episode` from the original tf.data pipeline. It operates on
    a single data record (episode).
    """

    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the transformation with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c

    def random_map(self, element: dict, rng: np.random.Generator) -> Any:
        """
        Processes a single raw episode from the data source.

        Args:
            element: A dictionary representing one record from the DataSource.
                     Expected to contain 'raw_video' (bytes), and
                     'sequence_length' (int).
            rng: A per-record random number generator provided by the Grain sampler.

        Returns:
            A processed video sequence as a NumPy array with shape
            (seq_len, image_h, image_w, image_c) and dtype float32.
        """
        # Step 1: Parse and reshape the raw video data (replaces _parse_tfrecord_fn)
        # We assume the element from the data source is a dict.
        video_shape = (
            element["sequence_length"],
            self.image_h,
            self.image_w,
            self.image_c,
        )
        episode_tensor = np.frombuffer(element["raw_video"], dtype=np.uint8)
        episode_tensor = episode_tensor.reshape(video_shape)

        # Step 2: Randomly slice and normalize (replaces _tf_process_episode)
        current_episode_len = episode_tensor.shape[0]
        if current_episode_len < self.seq_len:
             raise ValueError(f"An episode has length {current_episode_len}, which is "
                              f"shorter than the requested sequence length {self.seq_len}.")
        
        max_start_idx = current_episode_len - self.seq_len
        
        # Use the provided Grain RNG for deterministic randomness.
        start_idx = rng.integers(0, max_start_idx + 1)

        seq = episode_tensor[start_idx : start_idx + self.seq_len]

        # Normalize to [0, 1]
        processed_sequence = seq.astype(np.float32) / 255.0

        return processed_sequence


def get_dataloader_grain(
    array_record_paths: list[str],  # List of ArrayRecord file paths
    seq_len: int,
    global_batch_size: int,
    image_h: int,
    image_w: int,
    image_c: int,
    # shuffle_buffer_size is not needed; Grain shuffles all indices globally.
    num_workers: int = 4,  # Replaces num_parallel_calls
    seed: int = 42,
):
    """
    Creates a data loading pipeline using Google Grain.

    This function sets up a data source, a sampler for sharding/shuffling,
    and a sequence of operations (transformations) to be applied to the data.
    """
    if not array_record_paths:
        raise ValueError("array_record_paths list cannot be empty.")

    process_id = jax.process_index()
    num_processes = jax.process_count()

    if global_batch_size % num_processes != 0:
        raise ValueError(
            f"Global batch size {global_batch_size} must be divisible by "
            f"the number of JAX processes {num_processes} for proper sharding."
        )
    per_process_batch_size = global_batch_size // num_processes

    # 1. DataSource: Reads records from ArrayRecord files.
    # This replaces `tf.data.TFRecordDataset`.
    source = grain.sources.ArrayRecordDataSource(array_record_paths)
    
    # 2. Sampler: Defines the order of access, sharding, shuffling, and epochs.
    # This replaces `dataset.shard`, `dataset.shuffle`, and `dataset.repeat`.
    shard_options = grain.sharding.ShardOptions(
        shard_index=process_id, shard_count=num_processes, drop_remainder=True
    )

    sampler = grain.samplers.IndexSampler(
        num_records=len(source),
        shard_options=shard_options,
        seed=seed,
        num_epochs=None,  # `None` means repeat indefinitely, like `tf.data.repeat(None)`.
        shuffle=True,  # Enables global shuffling of all record indices.
    )

    # 3. Operations: A list of transformations to apply to each record.
    # This replaces the sequence of `.map()` calls.
    operations = [
        ProcessEpisodeAndSlice(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        grain.transforms.Batch(batch_size=per_process_batch_size, drop_remainder=True),
    ]

    # 4. DataLoader: The main entry point that orchestrates the pipeline.
    # It manages worker processes for parallel data loading and transformation,
    # replacing `num_parallel_calls` and `.prefetch()`.
    dataloader = grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=num_workers,
    )

    # 5. Return an iterator, which is the standard way to consume a Grain DataLoader.
    return iter(dataloader)