def get_dataset(
    dataset: str,
    global_batch_size: int,
    rng: np.ndarray,
    train_preprocessing_fn: Optional[Callable[[Any], Any]] = None,
    eval_preprocessing_fn: Optional[Callable[[Any], Any]] = None,
    num_epochs: Optional[int] = None,
    filter_fn: Optional[Callable[[Any], Any]] = None,
    **kwargs,