def PadToLength(  # pylint: disable=invalid-name
    len_map=None, pad_value=0, multiple=False):
  """Pads the values to lengths given in `len_map'.

  len_map contains a dictionary of example keys to dimension sizes.

  Args:
    len_map: dict of int to int, we pad examples to lengths
      given by the values of the dict. If multiple is True, the dimensions are
      padded to multiple of this value.
    pad_value: dict of int to int. The value gets applied to
      constant_values on numpy.pad per given dimension.
    multiple: boolean. If False, pads to the value of len_map. If True, pads to
      closest multiple of value of len_map.
  Returns:
    Function to pad examples to given lengths.
  """
  @debug_data_pipeline.debug_pipeline
  def _pad_to_length(generator):
    for example in generator:
      if isinstance(example, (list, tuple)):
        example = list(example)
        for key, value in len_map.items():
          array_length = example[key].shape[0]
          if multiple:
            padding_len = array_length - ((array_length // value) * value)
          else:
            padding_len = max([0, value-example[key].shape[0]])
          example[key] = np.pad(example[key],
                                pad_width=(0, padding_len),
                                mode='constant',
                                constant_values=pad_value[key])
        output = tuple(example)
      else:
        if not isinstance(example, np.ndarray):
          raise ValueError(f'example isn\'t nparray, but should be: {example}')
        array_length = example.shape[0]
        if multiple:
          padding_len = (
              array_length - ((array_length // len_map[0]) * len_map[0]))
        else:
          padding_len = max(0, len_map[0] - array_length)
        output = np.pad(example,
                        pad_width=(0, padding_len),
                        mode='constant',
                        constant_values=pad_value[0])
      yield output
  if len_map is None:
    raise ValueError('len_map parameter should be provided.')
  return _pad_to_length
Pads the values to lengths given in `len_map'.

  len_map contains a dictionary of example keys to dimension sizes.

  Args:
    len_map: dict of int to int, we pad examples to lengths
      given by the values of the dict. If multiple is True, the dimensions are
      padded to multiple of this value.
    pad_value: dict of int to int. The value gets applied to
      constant_values on numpy.pad per given dimension.
    multiple: boolean. If False, pads to the value of len_map. If True, pads to
      closest multiple of value of len_map.
  Returns:
    Function to pad examples to given lengths.
  
trax/data/inputs.py