def truncate_dataset_on_len(dataset,
                            training,
                            len_map=None,
                            truncate_on_eval=False):
  """Truncates features in an example to lengths given in `len_map`.

  Args:
    dataset: `tf.data.Dataset` the dataset to filter.
    training: bool, true if we are in training mode.
    len_map: optional dict of str to int, we truncate examples where a feature's
      size is beyond the max. Ex: {'inputs': 512, 'targets': 64} will truncate
        examples to be within those bounds.
    truncate_on_eval: bool if true, we will truncate in eval mode also.

  Returns:
    a filtered `tf.data.Dataset`.
  """
  if (len_map is None) or (not training and not truncate_on_eval):
    return dataset

  assert isinstance(len_map, dict)

  def truncate_example(x):
    for key, max_len in len_map.items():
      x_len = tf.shape(x[key])[0]
      if x_len > max_len:
        x[key] = x[key][:max_len, ...]
    return x

  return dataset.map(truncate_example)
Truncates features in an example to lengths given in `len_map`.

  Args:
    dataset: `tf.data.Dataset` the dataset to filter.
    training: bool, true if we are in training mode.
    len_map: optional dict of str to int, we truncate examples where a feature's
      size is beyond the max. Ex: {'inputs': 512, 'targets': 64} will truncate
        examples to be within those bounds.
    truncate_on_eval: bool if true, we will truncate in eval mode also.

  Returns:
    a filtered `tf.data.Dataset`.
  
trax/data/tf_inputs.py