def _train_and_eval_streams(dataset, data_dir, preprocess_fn,
                            bare_preprocess_fn, shuffle_buffer_size,
                            eval_holdout_size, input_name, target_name):
  """Return train and eval batches with input name and shape."""
  (train_data, eval_data,
   keys) = _train_and_eval_dataset(dataset, data_dir, eval_holdout_size)
  # If provided select input_name/target_name else fall back to keys if that is
  # available, else [None].
  input_names = ([input_name] if input_name is not None else
                 keys[0] if keys is not None else [None])
  target_names = ([target_name] if target_name is not None else
                  keys[1] if keys is not None else [None])

  train_batches = _shuffle_data(train_data, target_names, True,
                                shuffle_buffer_size, preprocess_fn,
                                bare_preprocess_fn)
  eval_batches = _shuffle_data(eval_data, target_names, False,
                               shuffle_buffer_size, preprocess_fn,
                               bare_preprocess_fn)
  return (train_batches, eval_batches, input_names[0])
Return train and eval batches with input name and shape.
trax/data/tf_inputs.py