def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
    """
    Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
    Shuffle batches if `shuffle` is `True`.
    """
    steps = len(dataset) // batch_size  # Skip incomplete batch.

    # We use `numpy.ndarray` to interact with `datasets.Dataset`, since using `jax.numpy.array` to index into a
    # dataset is significantly slow. Using JAX array at the 1st place is only to keep JAX's PRNGs generation
    # mechanism, which works differently from NumPy/SciPy.
    if shuffle:
        batch_idx = jax.random.permutation(rng, len(dataset))
        batch_idx = np.asarray(batch_idx)
    else:
        batch_idx = np.arange(len(dataset))

    for idx in range(steps):
        start_idx = batch_size * idx
        end_idx = batch_size * (idx + 1)

        selected_indices = batch_idx[start_idx:end_idx]
        batch = dataset[selected_indices]
        batch = shard(batch)

        yield batch

    Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
    Shuffle batches if `shuffle` is `True`.
    
dpkd/transformers/examples/flax/image-captioning/run_image_captioning_flax.py