def read_tf_records(batch_size, tf_records, num_repeats=1,
                    shuffle_records=True, shuffle_examples=True,
                    shuffle_buffer_size=None, interleave=True,
                    filter_amount=1.0):
  """
    Args:
        batch_size: batch size to return
        tf_records: a list of tf_record filenames
        num_repeats: how many times the data should be read (default: One)
        shuffle_records: whether to shuffle the order of files read
        shuffle_examples: whether to shuffle the tf.Examples
        shuffle_buffer_size: how big of a buffer to fill before shuffling.
        interleave: iwhether to interleave examples from multiple tf_records
        filter_amount: what fraction of records to keep

    Returns:
        a tf dataset of batched tensors
  """
  if shuffle_examples and not shuffle_buffer_size:
    raise ValueError('Must set shuffle buffer size if shuffling examples')

  tf_records = list(tf_records)
  if shuffle_records:
    random.shuffle(tf_records)
  record_list = tf.data.Dataset.from_tensor_slices(tf_records)

  # compression_type here must agree with write_tf_examples
  map_func = functools.partial(
      tf.data.TFRecordDataset,
      buffer_size=8 * 1024 * 1024,
      compression_type='ZLIB')

  if interleave:
    # cycle_length = how many tfrecord files are read in parallel
    # The idea is to shuffle both the order of the files being read,
    # and the examples being read from the files.
    dataset = record_list.apply(
        tf.data.experimental.parallel_interleave(
            map_func, cycle_length=64, sloppy=True))
  else:
    dataset = record_list.flat_map(map_func)

  if filter_amount < 1.0:
    dataset = dataset.filter(lambda _: tf.random_uniform([]) < filter_amount)

  dataset = dataset.repeat(num_repeats)
  if shuffle_examples:
    dataset = dataset.shuffle(buffer_size=shuffle_buffer_size)

  dataset = dataset.batch(batch_size)
  return dataset