def stream(which):
    """Create the stream, cache TF streams if needed."""
    if not cache:
      cache.append(
          _train_and_eval_streams(dataset_name, data_dir, preprocess_fn,
                                  bare_preprocess_fn, shuffle_buffer_size,
                                  eval_holdout_size, input_name, target_name))

    (train_ds, eval_ds, input_name_c) = cache[0]
    dataset = eval_ds if which == 'eval' else train_ds
    return dataset_to_stream(dataset, input_name_c)
Create the stream, cache TF streams if needed.
trax/data/tf_inputs.py