def dataset_size_for_mode(mode):
  """Returns the number of training examples in the input dataset."""
  if mode == 'test':
    return 50000
  elif mode == 'train':
    return 1281167
  elif mode == 'l2l_valid':
    return 50046
  elif mode == 'l2l_train':
    return 1281167 - 50046
  else:
    raise ValueError('Invalid mode: {!r}'.format(mode))