def tfds_split_for_mode(mode):
  """Return the TFDS split to use for a given input dataset."""
  if mode == 'test':
    # The labels for the real ImageNet test set were never released. So we
    # follow the standard (although admitted confusing) practice of obtain
    # "test set" accuracy numbers from the ImageNet validation.
    return 'validation'
  elif mode == 'train':
    return 'train'
  elif mode == 'l2l_valid':
    # To prevent overfitting to the test set, we  use a held-out portion of the
    # training set for model validation. We use the same number of examples here
    # as we did for the TuNAS paper. However, the validation set contains
    # different examples than it did in the paper, and we should expect small
    # deviations from the reported results for this reason.
    return 'train[:50046]'
  elif mode == 'l2l_train':
    return 'train[50046:]'
  else:
    raise ValueError('Invalid mode: {!r}'.format(mode))