nodes: 0, module; 1, function_definition; 2, function_name:sort_batch_by_length; 3, parameters; 4, block; 5, typed_parameter; 6, typed_parameter; 7, expression_statement; 8, if_statement; 9, expression_statement; 10, expression_statement; 11, expression_statement; 12, comment:# This is the equivalent of zipping with index, sorting by the original; 13, comment:# sequence lengths and returning the now sorted indices.; 14, expression_statement; 15, expression_statement; 16, return_statement; 17, identifier:tensor; 18, type; 19, identifier:sequence_lengths; 20, type; 21, comment:"""
      Sort a batch first tensor by some specified lengths.
  
      Parameters
      ----------
      tensor : torch.FloatTensor, required.
          A batch first Pytorch tensor.
      sequence_lengths : torch.LongTensor, required.
          A tensor representing the lengths of some dimension of the tensor which
          we want to sort by.
  
      Returns
      -------
      sorted_tensor : torch.FloatTensor
          The original tensor sorted along the batch dimension with respect to sequence_lengths.
      sorted_sequence_lengths : torch.LongTensor
          The original sequence_lengths sorted by decreasing size.
      restoration_indices : torch.LongTensor
          Indices into the sorted_tensor such that
          ``sorted_tensor.index_select(0, restoration_indices) == original_tensor``
      permutation_index : torch.LongTensor
          The indices used to sort the tensor. This is useful if you want to sort many
          tensors using the same ordering.
      """; 22, boolean_operator; 23, block; 24, assignment; 25, assignment; 26, assignment; 27, assignment; 28, assignment; 29, expression_list; 30, attribute; 31, attribute; 32, not_operator; 33, not_operator; 34, raise_statement; 35, pattern_list; 36, call; 37, identifier:sorted_tensor; 38, call; 39, identifier:index_range; 40, call; 41, pattern_list; 42, call; 43, identifier:restoration_indices; 44, call; 45, identifier:sorted_tensor; 46, identifier:sorted_sequence_lengths; 47, identifier:restoration_indices; 48, identifier:permutation_index; 49, identifier:torch; 50, identifier:Tensor; 51, identifier:torch; 52, identifier:Tensor; 53, call; 54, call; 55, call; 56, identifier:sorted_sequence_lengths; 57, identifier:permutation_index; 58, attribute; 59, argument_list; 60, attribute; 61, argument_list; 62, attribute; 63, argument_list; 64, identifier:_; 65, identifier:reverse_mapping; 66, attribute; 67, argument_list; 68, attribute; 69, argument_list; 70, identifier:isinstance; 71, argument_list; 72, identifier:isinstance; 73, argument_list; 74, identifier:ConfigurationError; 75, argument_list; 76, identifier:sequence_lengths; 77, identifier:sort; 78, integer:0; 79, keyword_argument; 80, identifier:tensor; 81, identifier:index_select; 82, integer:0; 83, identifier:permutation_index; 84, identifier:torch; 85, identifier:arange; 86, integer:0; 87, call; 88, keyword_argument; 89, identifier:permutation_index; 90, identifier:sort; 91, integer:0; 92, keyword_argument; 93, identifier:index_range; 94, identifier:index_select; 95, integer:0; 96, identifier:reverse_mapping; 97, identifier:tensor; 98, attribute; 99, identifier:sequence_lengths; 100, attribute; 101, string:"Both the tensor and sequence lengths must be torch.Tensors."; 102, identifier:descending; 103, True; 104, identifier:len; 105, argument_list; 106, identifier:device; 107, attribute; 108, identifier:descending; 109, False; 110, identifier:torch; 111, identifier:Tensor; 112, identifier:torch; 113, identifier:Tensor; 114, identifier:sequence_lengths; 115, identifier:sequence_lengths; 116, identifier:device
code: def sort_batch_by_length(tensor: torch.Tensor, sequence_lengths: torch.Tensor):
      """
      Sort a batch first tensor by some specified lengths.
  
      Parameters
      ----------
      tensor : torch.FloatTensor, required.
          A batch first Pytorch tensor.
      sequence_lengths : torch.LongTensor, required.
          A tensor representing the lengths of some dimension of the tensor which
          we want to sort by.
  
      Returns
      -------
      sorted_tensor : torch.FloatTensor
          The original tensor sorted along the batch dimension with respect to sequence_lengths.
      sorted_sequence_lengths : torch.LongTensor
          The original sequence_lengths sorted by decreasing size.
      restoration_indices : torch.LongTensor
          Indices into the sorted_tensor such that
          ``sorted_tensor.index_select(0, restoration_indices) == original_tensor``
      permutation_index : torch.LongTensor
          The indices used to sort the tensor. This is useful if you want to sort many
          tensors using the same ordering.
      """
  
      if not isinstance(tensor, torch.Tensor) or not isinstance(sequence_lengths, torch.Tensor):
          raise ConfigurationError("Both the tensor and sequence lengths must be torch.Tensors.")
  
      sorted_sequence_lengths, permutation_index = sequence_lengths.sort(0, descending=True)
      sorted_tensor = tensor.index_select(0, permutation_index)
  
      index_range = torch.arange(0, len(sequence_lengths), device=sequence_lengths.device)
      # This is the equivalent of zipping with index, sorting by the original
      # sequence lengths and returning the now sorted indices.
      _, reverse_mapping = permutation_index.sort(0, descending=False)
      restoration_indices = index_range.index_select(0, reverse_mapping)
      return sorted_tensor, sorted_sequence_lengths, restoration_indices, permutation_index