def split_dataset(
    dataset,
    num_per_split,
    num_splits = 3,
):
  """Helper to split a dataset for cross validaton.

  The first num_splits-1 datasets contain num_per_split examples.

  The last dataset contains the remaining number of examples.
  This often used to split a training set into more validation sets:
  e.g. train_old --> [valid_inner, valid_outer, train]

  Args:
    dataset: name of tfds dataset
    num_per_split: number of examples to have for each of the split off dataset.
    num_splits: number of splits to create.

  Returns:
    A list of the requested datasets.
  """
  new_datasets = []

  # make the first n_split-1 splits containing num_per_split examples
  for i in range(num_splits - 1):
    new_datasets.append(dataset.skip(num_per_split * i).take(num_per_split))
  # The remainder of the dataset
  new_datasets.append(dataset.skip(num_per_split * (num_splits - 1)))

  return new_datasets