nodes: 0, module; 1, function_definition; 2, function_name:sort_and_run_forward; 3, parameters; 4, block; 5, identifier:self; 6, typed_parameter; 7, typed_parameter; 8, typed_parameter; 9, typed_default_parameter; 10, expression_statement; 11, comment:# In some circumstances you may have sequences of zero length. ``pack_padded_sequence``; 12, comment:# requires all sequence lengths to be > 0, so remove sequences of zero length before; 13, comment:# calling self._module, then fill with zeros.; 14, comment:# First count how many sequences are empty.; 15, expression_statement; 16, expression_statement; 17, expression_statement; 18, expression_statement; 19, comment:# Now create a PackedSequence with only the non-empty, sorted sequences.; 20, expression_statement; 21, comment:# Prepare the initial states.; 22, if_statement; 23, comment:# Actually call the module on the sorted PackedSequence.; 24, expression_statement; 25, return_statement; 26, identifier:module; 27, type; 28, identifier:inputs; 29, type; 30, identifier:mask; 31, type; 32, identifier:hidden_state; 33, type; 34, None; 35, comment:"""
          This function exists because Pytorch RNNs require that their inputs be sorted
          before being passed as input. As all of our Seq2xxxEncoders use this functionality,
          it is provided in a base class. This method can be called on any module which
          takes as input a ``PackedSequence`` and some ``hidden_state``, which can either be a
          tuple of tensors or a tensor.
  
          As all of our Seq2xxxEncoders have different return types, we return `sorted`
          outputs from the module, which is called directly. Additionally, we return the
          indices into the batch dimension required to restore the tensor to it's correct,
          unsorted order and the number of valid batch elements (i.e the number of elements
          in the batch which are not completely masked). This un-sorting and re-padding
          of the module outputs is left to the subclasses because their outputs have different
          types and handling them smoothly here is difficult.
  
          Parameters
          ----------
          module : ``Callable[[PackedSequence, Optional[RnnState]],
                              Tuple[Union[PackedSequence, torch.Tensor], RnnState]]``, required.
              A function to run on the inputs. In most cases, this is a ``torch.nn.Module``.
          inputs : ``torch.Tensor``, required.
              A tensor of shape ``(batch_size, sequence_length, embedding_size)`` representing
              the inputs to the Encoder.
          mask : ``torch.Tensor``, required.
              A tensor of shape ``(batch_size, sequence_length)``, representing masked and
              non-masked elements of the sequence for each element in the batch.
          hidden_state : ``Optional[RnnState]``, (default = None).
              A single tensor of shape (num_layers, batch_size, hidden_size) representing the
              state of an RNN with or a tuple of
              tensors of shapes (num_layers, batch_size, hidden_size) and
              (num_layers, batch_size, memory_size), representing the hidden state and memory
              state of an LSTM-like RNN.
  
          Returns
          -------
          module_output : ``Union[torch.Tensor, PackedSequence]``.
              A Tensor or PackedSequence representing the output of the Pytorch Module.
              The batch size dimension will be equal to ``num_valid``, as sequences of zero
              length are clipped off before the module is called, as Pytorch cannot handle
              zero length sequences.
          final_states : ``Optional[RnnState]``
              A Tensor representing the hidden state of the Pytorch Module. This can either
              be a single tensor of shape (num_layers, num_valid, hidden_size), for instance in
              the case of a GRU, or a tuple of tensors, such as those required for an LSTM.
          restoration_indices : ``torch.LongTensor``
              A tensor of shape ``(batch_size,)``, describing the re-indexing required to transform
              the outputs back to their original batch order.
          """; 36, assignment; 37, assignment; 38, assignment; 39, assignment; 40, assignment; 41, not_operator; 42, block; 43, else_clause; 44, assignment; 45, expression_list; 46, generic_type; 47, attribute; 48, attribute; 49, generic_type; 50, identifier:batch_size; 51, call; 52, identifier:num_valid; 53, call; 54, identifier:sequence_lengths; 55, call; 56, pattern_list; 57, line_continuation:\; 58, call; 59, identifier:packed_sequence_input; 60, call; 61, attribute; 62, if_statement; 63, block; 64, pattern_list; 65, call; 66, identifier:module_output; 67, identifier:final_states; 68, identifier:restoration_indices; 69, identifier:Callable; 70, type_parameter; 71, identifier:torch; 72, identifier:Tensor; 73, identifier:torch; 74, identifier:Tensor; 75, identifier:Optional; 76, type_parameter; 77, attribute; 78, argument_list; 79, attribute; 80, argument_list; 81, identifier:get_lengths_from_binary_sequence_mask; 82, argument_list; 83, identifier:sorted_inputs; 84, identifier:sorted_sequence_lengths; 85, identifier:restoration_indices; 86, identifier:sorting_indices; 87, identifier:sort_batch_by_length; 88, argument_list; 89, identifier:pack_padded_sequence; 90, argument_list; 91, identifier:self; 92, identifier:stateful; 93, comparison_operator:hidden_state is None; 94, block; 95, elif_clause; 96, else_clause; 97, expression_statement; 98, identifier:module_output; 99, identifier:final_states; 100, identifier:module; 101, argument_list; 102, type; 103, type; 104, type; 105, identifier:mask; 106, identifier:size; 107, integer:0; 108, call; 109, identifier:item; 110, identifier:mask; 111, identifier:inputs; 112, identifier:sequence_lengths; 113, subscript; 114, call; 115, keyword_argument; 116, identifier:hidden_state; 117, None; 118, expression_statement; 119, call; 120, block; 121, block; 122, assignment; 123, identifier:packed_sequence_input; 124, identifier:initial_states; 125, list; 126, generic_type; 127, identifier:RnnState; 128, attribute; 129, argument_list; 130, identifier:sorted_inputs; 131, slice; 132, slice; 133, slice; 134, attribute; 135, argument_list; 136, identifier:batch_first; 137, True; 138, assignment; 139, identifier:isinstance; 140, argument_list; 141, expression_statement; 142, expression_statement; 143, identifier:initial_states; 144, call; 145, identifier:PackedSequence; 146, subscript; 147, identifier:Tuple; 148, type_parameter; 149, call; 150, identifier:int; 151, identifier:num_valid; 152, attribute; 153, identifier:tolist; 154, identifier:initial_states; 155, identifier:hidden_state; 156, identifier:hidden_state; 157, identifier:tuple; 158, assignment; 159, assignment; 160, attribute; 161, argument_list; 162, identifier:Optional; 163, identifier:RnnState; 164, type; 165, type; 166, attribute; 167, argument_list; 168, subscript; 169, identifier:data; 170, identifier:initial_states; 171, list_comprehension; 172, identifier:initial_states; 173, call; 174, identifier:self; 175, identifier:_get_initial_states; 176, identifier:batch_size; 177, identifier:num_valid; 178, identifier:sorting_indices; 179, generic_type; 180, identifier:RnnState; 181, identifier:torch; 182, identifier:sum; 183, subscript; 184, identifier:sorted_sequence_lengths; 185, slice; 186, call; 187, for_in_clause; 188, attribute; 189, argument_list; 190, identifier:Union; 191, type_parameter; 192, identifier:mask; 193, slice; 194, integer:0; 195, identifier:num_valid; 196, attribute; 197, argument_list; 198, identifier:state; 199, identifier:hidden_state; 200, subscript; 201, identifier:contiguous; 202, type; 203, type; 204, subscript; 205, identifier:contiguous; 206, call; 207, slice; 208, slice; 209, slice; 210, identifier:PackedSequence; 211, attribute; 212, call; 213, slice; 214, slice; 215, slice; 216, attribute; 217, argument_list; 218, identifier:num_valid; 219, identifier:torch; 220, identifier:Tensor; 221, attribute; 222, argument_list; 223, identifier:num_valid; 224, identifier:hidden_state; 225, identifier:index_select; 226, integer:1; 227, identifier:sorting_indices; 228, identifier:state; 229, identifier:index_select; 230, integer:1; 231, identifier:sorting_indices
code: def sort_and_run_forward(self,
                               module: Callable[[PackedSequence, Optional[RnnState]],
                                                Tuple[Union[PackedSequence, torch.Tensor], RnnState]],
                               inputs: torch.Tensor,
                               mask: torch.Tensor,
                               hidden_state: Optional[RnnState] = None):
          """
          This function exists because Pytorch RNNs require that their inputs be sorted
          before being passed as input. As all of our Seq2xxxEncoders use this functionality,
          it is provided in a base class. This method can be called on any module which
          takes as input a ``PackedSequence`` and some ``hidden_state``, which can either be a
          tuple of tensors or a tensor.
  
          As all of our Seq2xxxEncoders have different return types, we return `sorted`
          outputs from the module, which is called directly. Additionally, we return the
          indices into the batch dimension required to restore the tensor to it's correct,
          unsorted order and the number of valid batch elements (i.e the number of elements
          in the batch which are not completely masked). This un-sorting and re-padding
          of the module outputs is left to the subclasses because their outputs have different
          types and handling them smoothly here is difficult.
  
          Parameters
          ----------
          module : ``Callable[[PackedSequence, Optional[RnnState]],
                              Tuple[Union[PackedSequence, torch.Tensor], RnnState]]``, required.
              A function to run on the inputs. In most cases, this is a ``torch.nn.Module``.
          inputs : ``torch.Tensor``, required.
              A tensor of shape ``(batch_size, sequence_length, embedding_size)`` representing
              the inputs to the Encoder.
          mask : ``torch.Tensor``, required.
              A tensor of shape ``(batch_size, sequence_length)``, representing masked and
              non-masked elements of the sequence for each element in the batch.
          hidden_state : ``Optional[RnnState]``, (default = None).
              A single tensor of shape (num_layers, batch_size, hidden_size) representing the
              state of an RNN with or a tuple of
              tensors of shapes (num_layers, batch_size, hidden_size) and
              (num_layers, batch_size, memory_size), representing the hidden state and memory
              state of an LSTM-like RNN.
  
          Returns
          -------
          module_output : ``Union[torch.Tensor, PackedSequence]``.
              A Tensor or PackedSequence representing the output of the Pytorch Module.
              The batch size dimension will be equal to ``num_valid``, as sequences of zero
              length are clipped off before the module is called, as Pytorch cannot handle
              zero length sequences.
          final_states : ``Optional[RnnState]``
              A Tensor representing the hidden state of the Pytorch Module. This can either
              be a single tensor of shape (num_layers, num_valid, hidden_size), for instance in
              the case of a GRU, or a tuple of tensors, such as those required for an LSTM.
          restoration_indices : ``torch.LongTensor``
              A tensor of shape ``(batch_size,)``, describing the re-indexing required to transform
              the outputs back to their original batch order.
          """
          # In some circumstances you may have sequences of zero length. ``pack_padded_sequence``
          # requires all sequence lengths to be > 0, so remove sequences of zero length before
          # calling self._module, then fill with zeros.
  
          # First count how many sequences are empty.
          batch_size = mask.size(0)
          num_valid = torch.sum(mask[:, 0]).int().item()
  
          sequence_lengths = get_lengths_from_binary_sequence_mask(mask)
          sorted_inputs, sorted_sequence_lengths, restoration_indices, sorting_indices =\
              sort_batch_by_length(inputs, sequence_lengths)
  
          # Now create a PackedSequence with only the non-empty, sorted sequences.
          packed_sequence_input = pack_padded_sequence(sorted_inputs[:num_valid, :, :],
                                                       sorted_sequence_lengths[:num_valid].data.tolist(),
                                                       batch_first=True)
          # Prepare the initial states.
          if not self.stateful:
              if hidden_state is None:
                  initial_states = hidden_state
              elif isinstance(hidden_state, tuple):
                  initial_states = [state.index_select(1, sorting_indices)[:, :num_valid, :].contiguous()
                                    for state in hidden_state]
              else:
                  initial_states = hidden_state.index_select(1, sorting_indices)[:, :num_valid, :].contiguous()
  
          else:
              initial_states = self._get_initial_states(batch_size, num_valid, sorting_indices)
  
          # Actually call the module on the sorted PackedSequence.
          module_output, final_states = module(packed_sequence_input, initial_states)
  
          return module_output, final_states, restoration_indices