def _prepare_dataset(
    dataset: tf.data.Dataset,
    global_batch_size: int,
    shuffle: bool,
    rng: np.ndarray,
    preprocess_fn: Optional[Callable[[Any], Any]] = None,
    num_epochs: Optional[int] = None,
    filter_fn: Optional[Callable[[Any], Any]] = None) -> tf.data.Dataset:
  """Batches, shuffles, prefetches and preprocesses a dataset.

  Args:
    dataset: The dataset to prepare.
    global_batch_size: The global batch size to use.
    shuffle: Whether the shuffle the data on example level.
    rng: PRNG for seeding the shuffle operations.
    preprocess_fn: Preprocessing function that will be applied to every example.
    num_epochs: Number of epochs to repeat the dataset.
    filter_fn: Funtion that filters samples according to some criteria.

  Returns:
    The dataset.
  """
  if shuffle and rng is None:
    raise ValueError("Shuffling without RNG is not supported.")

  if global_batch_size % jax.host_count() != 0:
    raise ValueError(f"Batch size {global_batch_size} not divisible by number "
                     f"of hosts ({jax.host_count()}).")
  local_batch_size = global_batch_size // jax.host_count()
  batch_dims = [jax.local_device_count(), local_batch_size]

  # tf.data uses single integers as seed.
  if rng is not None:
    rng = rng[0]

  ds = dataset.repeat(num_epochs)
  if shuffle:
    ds = ds.shuffle(1024, seed=rng)

  if preprocess_fn is not None:
    ds = ds.map(preprocess_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)

  if filter_fn is not None:
    ds = ds.filter(filter_fn)

  for batch_size in reversed(batch_dims):
    ds = ds.batch(batch_size, drop_remainder=True)
  return ds.prefetch(tf.data.experimental.AUTOTUNE)