def TruncateToLength(len_map=None):  # pylint: disable=invalid-name
  """Returns a stream function that resizes items as specified by ``len_map``.

  Args:
    len_map: Dictionary that specifies maximum shapes for potentially multiple
        features per stream item. For example, given a stream of tokenized
        string pairs, one could enforce a maximum length of 256 tokens for each
        string by using ``len_map={0: (256,), 1: (256,)}``.
  """
  @debug_data_pipeline.debug_pipeline
  def _truncate_to_length(generator):
    for example in generator:
      if isinstance(example, np.ndarray):
        example = (example,)
      if isinstance(example, (list, tuple)):
        example = list(example)
        if len_map is not None:
          for key, max_len in len_map.items():
            example_len = example[key].shape
            if example_len > max_len:
              example[key] = np.resize(example[key], max_len)
        output = tuple(example)
      else:
        output = None
        raise ValueError(f'Unknown example type: {example}')
      yield output

  return _truncate_to_length
Returns a stream function that resizes items as specified by ``len_map``.

  Args:
    len_map: Dictionary that specifies maximum shapes for potentially multiple
        features per stream item. For example, given a stream of tokenized
        string pairs, one could enforce a maximum length of 256 tokens for each
        string by using ``len_map={0: (256,), 1: (256,)}``.
  
trax/data/inputs.py