(  # 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)
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.
Creates a data source from TensorFlow dataset ``dataset_name``.
Creates a data source from TensorFlow dataset dataset_name .
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
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 ) : 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