def blockwise_data_loader(
        rng: jax.random.PRNGKey,
        ds: Dataset,
        block_size: int,
        batch_size: int,
        shuffle: bool = False,
        keep_in_memory: bool = False,
        split: str = "",
    ):
        """
        Wrap the simple `data_loader` in a block-wise way if `block_size` > 0, else it's the same as `data_loader`.

        If `block_size` > 0, it requires `ds` to have a column that gives image paths in order to perform image
        processing (with the column name being specified by `image_column`). The tokenization should be done before
        training in this case.
        """

        # 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:
            indices = jax.random.permutation(rng, len(ds))
            indices = np.asarray(indices)
        else:
            indices = np.arange(len(ds))

        _block_size = len(ds) if not block_size else block_size

        steps_per_block = _block_size // batch_size
        num_examples = len(ds)
        steps = num_examples // batch_size
        num_splits = steps // steps_per_block + int(steps % steps_per_block > 0)

        for idx in range(num_splits):
            if not block_size:
                _ds = ds
            else:
                start_idx = block_size * idx
                end_idx = block_size * (idx + 1)

                selected_indices = indices[start_idx:end_idx]

                _ds = ds.select(selected_indices)

                _ds = _ds.map(
                    image_processing_fn,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=[image_column],
                    load_from_cache_file=not data_args.overwrite_cache,
                    features=features,
                    keep_in_memory=keep_in_memory,
                    # The images are already checked either in `.filter()` or in `preprocess_fn()`
                    fn_kwargs={"check_image": False},
                    desc=f"Running image processing on {split} dataset".replace("  ", " "),
                )
                _ds = _ds.with_format("numpy")

            # No need to shuffle here
            loader = data_loader(rng, _ds, batch_size=batch_size, shuffle=False)

            for batch in loader:
                yield batch

        Wrap the simple `data_loader` in a block-wise way if `block_size` > 0, else it's the same as `data_loader`.

        If `block_size` > 0, it requires `ds` to have a column that gives image paths in order to perform image
        processing (with the column name being specified by `image_column`). The tokenization should be done before
        training in this case.
        
dpkd/transformers/examples/flax/image-captioning/run_image_captioning_flax.py