examples/research_projects/performer/run_mlm_performer.py/generate_batch_splits
def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
    nb_samples = len(samples_idx)
    samples_to_remove = nb_samples % batch_size

    if samples_to_remove != 0:
        samples_idx = samples_idx[:-samples_to_remove]
    sections_split = nb_samples // batch_size
    batch_idx = np.split(samples_idx, sections_split)
    return batch_idx