def ConcatenateToLMInput(pad_to_length=None):  # pylint: disable=invalid-name
  """Prepares the input needed for training of Language Models.

  Each example needs to contain two elements (input and target).
  Input is concatenated to target and, if pad_to_length is given, padded to
  length provided.
  The loss_weights indicates only the target, without input nor padding.

  Args:
    pad_to_length: int, total length of padding of input and target arrays.
  Returns:
    Function to return input for a LM.
  """
  @debug_data_pipeline.debug_pipeline
  def _concatenate_to_lm_input(generator):
    for example in generator:
      if isinstance(example, (list, tuple)) and (len(example) == 2):
        concatenated = np.concatenate((example[0], example[1]), axis=-1)
        loss_weights = np.concatenate((np.zeros_like(example[0]),
                                       np.ones_like(example[1])))
        if pad_to_length is not None:
          padding_len = pad_to_length - (
              example[0].shape[0] + example[1].shape[0])
          if padding_len < 0:
            raise ValueError(
                'Example lengths '
                f'({example[0].shape[0]}, {example[1].shape[0]}) '
                f'longer than pad_to_length ({pad_to_length}).')
          loss_weights = np.pad(loss_weights, (0, padding_len), 'constant')
          concatenated = np.pad(concatenated, (0, padding_len), 'constant')
        output = (concatenated, concatenated, loss_weights)
      elif isinstance(example, (list, tuple)) and (len(example) == 1):
        # Make x into (x, x)
        output = (example[0], example[0])
      elif isinstance(example, np.ndarray):
        # Make x into (x, x)
        output = (example, example)
      else:
        output = None
        raise ValueError(f'Unknown input to ConcatenateToLMInput: {example}')
      yield output

  return _concatenate_to_lm_input
Prepares the input needed for training of Language Models.

  Each example needs to contain two elements (input and target).
  Input is concatenated to target and, if pad_to_length is given, padded to
  length provided.
  The loss_weights indicates only the target, without input nor padding.

  Args:
    pad_to_length: int, total length of padding of input and target arrays.
  Returns:
    Function to return input for a LM.
  
trax/data/inputs.py