f METHOD_NAME(self,
                     data_rng: jax.random.PRNGKey,
                     split: str,
                     data_dir: str,
                     global_batch_size: int,
                     num_batches: Optional[int] = None,
                     repeat_final_dataset: bool = False):