def TFDS(  # pylint: disable=invalid-name
    dataset_name,
    data_dir=None,
    tfds_preprocess_fn=None,
    keys=None,
    train=True,
    use_alt_eval=False,
    shuffle_train=True,
    host_id=None,
    n_hosts=None,
    eval_holdout_size=0):
  """Creates a data source from TensorFlow dataset ``dataset_name``.

  Args:
    dataset_name: Name of the dataset, as registered in TensorFlow datasets
      (e.g., ``'glue/mnli'``).
    data_dir: Directory where the data is located.
    tfds_preprocess_fn: If specified, function that applies to items in raw
      dataset (before selecting specific features).
    keys: Tuple of dataset-specific strings that select features from the
      dataset.
    train: If True, select the training split from the dataset; else select an
      eval split.
    use_alt_eval: If True, and if ``train`` is False, select the dataset's
      alternate eval split if it has one (or fall back to the dataset's only
      eval split). This currently affects only the `glue/mnli` dataset.
    shuffle_train: If True, have TensorFlow pre-shuffle the training data; else
      receive training data in deterministic sequence.
    host_id: Integer id used for tracking data subsplits, in cases where
      ``n_hosts`` > 1.
    n_hosts: If greater than 1, prepare data subsplits for the given number of
      hosts.
    eval_holdout_size: If greater than 0, specifies a fraction of training data
      to siphon off and use as eval data, in place of an separate eval split.

  Returns:
    A function `f` for use as a training or eval data source; call `f()` to get
    an iterator of data samples.
  """
  data_dir = download_and_prepare(dataset_name, data_dir)

  host_id = jax.process_index() if host_id is None else host_id
  n_hosts = n_hosts or jax.host_count()
  if n_hosts > 1:
    subsplit = (host_id / n_hosts, (host_id + 1) / n_hosts)
  else:
    subsplit = None
  train_data, eval_data, _ = (
      _train_and_eval_dataset(dataset_name,
                              data_dir,
                              eval_holdout_size,
                              train_shuffle_files=shuffle_train,
                              use_alt_eval=use_alt_eval,
                              subsplit=subsplit))
  dataset = train_data if train else eval_data
  dataset = dataset if tfds_preprocess_fn is None else tfds_preprocess_fn(
      dataset)

  def select_from(example):
    return tuple(example[k] for k in keys)

  dataset = dataset.map(select_from)
  dataset = dataset.repeat()

  def gen(generator=None):
    del generator
    for example in fastmath.dataset_as_numpy(dataset):
      yield example

  return gen
Creates a data source from TensorFlow dataset ``dataset_name``.

  Args:
    dataset_name: Name of the dataset, as registered in TensorFlow datasets
      (e.g., ``'glue/mnli'``).
    data_dir: Directory where the data is located.
    tfds_preprocess_fn: If specified, function that applies to items in raw
      dataset (before selecting specific features).
    keys: Tuple of dataset-specific strings that select features from the
      dataset.
    train: If True, select the training split from the dataset; else select an
      eval split.
    use_alt_eval: If True, and if ``train`` is False, select the dataset's
      alternate eval split if it has one (or fall back to the dataset's only
      eval split). This currently affects only the `glue/mnli` dataset.
    shuffle_train: If True, have TensorFlow pre-shuffle the training data; else
      receive training data in deterministic sequence.
    host_id: Integer id used for tracking data subsplits, in cases where
      ``n_hosts`` > 1.
    n_hosts: If greater than 1, prepare data subsplits for the given number of
      hosts.
    eval_holdout_size: If greater than 0, specifies a fraction of training data
      to siphon off and use as eval data, in place of an separate eval split.

  Returns:
    A function `f` for use as a training or eval data source; call `f()` to get
    an iterator of data samples.
  
trax/data/tf_inputs.py