def _shuffle_data(dataset, target_names, training, shuffle_buffer_size,
                  preprocess_fn, bare_preprocess_fn):
  """Shuffle the given dataset and run pre-processing."""

  def append_targets(example):
    """Append targets to the example dictionary. Needed for Keras."""
    if len(target_names) == 1:
      return (example, example[target_names[0]])
    targets = {}
    for name in target_names:
      targets[name] = example[name]
    return (example, targets)

  # `bare_preprocess_fn` is called before appending targets etc.
  if bare_preprocess_fn is not None:
    dataset = bare_preprocess_fn(dataset, training)
  dataset = dataset.map(append_targets)
  # TODO(pkozakowski): Repeat both the training and evaluation set, so we don't
  # have incomplete batches during evaluation. This will be a problem when we
  # add an option to evaluate on the whole dataset, then we'll need to think of
  # a different solution.
  dataset = dataset.repeat()
  if training:
    # Skip a random fraction at the beginning of the stream.  The skip is
    # essential for synchronous highly-parallel training to avoid multiple
    # replicas reading the same data in lock-step.
    dataset = dataset.skip(random.randint(0, int(_MAX_SKIP_EXAMPLES)))
  dataset = preprocess_fn(dataset, training)
  dataset = dataset.shuffle(shuffle_buffer_size)
  return dataset.prefetch(8)
Shuffle the given dataset and run pre-processing.
trax/data/tf_inputs.py