nodes: 0, module; 1, function_definition; 2, function_name:_get_initial_states; 3, parameters; 4, type; 5, block; 6, identifier:self; 7, typed_parameter; 8, typed_parameter; 9, typed_parameter; 10, generic_type; 11, expression_statement; 12, comment:# We don't know the state sizes the first time calling forward,; 13, comment:# so we let the module define what it's initial hidden state looks like.; 14, if_statement; 15, comment:# Otherwise, we have some previous states.; 16, if_statement; 17, comment:# At this point, our states are of shape (num_layers, batch_size, hidden_size).; 18, comment:# However, the encoder uses sorted sequences and additionally removes elements; 19, comment:# of the batch which are fully padded. We need the states to match up to these; 20, comment:# sorted and filtered sequences, so we do that in the next two blocks before; 21, comment:# returning the state/s.; 22, if_statement; 23, identifier:batch_size; 24, type; 25, identifier:num_valid; 26, type; 27, identifier:sorting_indices; 28, type; 29, identifier:Optional; 30, type_parameter; 31, comment:"""
          Returns an initial state for use in an RNN. Additionally, this method handles
          the batch size changing across calls by mutating the state to append initial states
          for new elements in the batch. Finally, it also handles sorting the states
          with respect to the sequence lengths of elements in the batch and removing rows
          which are completely padded. Importantly, this `mutates` the state if the
          current batch size is larger than when it was previously called.
  
          Parameters
          ----------
          batch_size : ``int``, required.
              The batch size can change size across calls to stateful RNNs, so we need
              to know if we need to expand or shrink the states before returning them.
              Expanded states will be set to zero.
          num_valid : ``int``, required.
              The batch may contain completely padded sequences which get removed before
              the sequence is passed through the encoder. We also need to clip these off
              of the state too.
          sorting_indices ``torch.LongTensor``, required.
              Pytorch RNNs take sequences sorted by length. When we return the states to be
              used for a given call to ``module.forward``, we need the states to match up to
              the sorted sequences, so before returning them, we sort the states using the
              same indices used to sort the sequences.
  
          Returns
          -------
          This method has a complex return type because it has to deal with the first time it
          is called, when it has no state, and the fact that types of RNN have heterogeneous
          states.
  
          If it is the first time the module has been called, it returns ``None``, regardless
          of the type of the ``Module``.
  
          Otherwise, for LSTMs, it returns a tuple of ``torch.Tensors`` with shape
          ``(num_layers, num_valid, state_size)`` and ``(num_layers, num_valid, memory_size)``
          respectively, or for GRUs, it returns a single ``torch.Tensor`` of shape
          ``(num_layers, num_valid, state_size)``.
          """; 32, comparison_operator:self._states is None; 33, block; 34, comparison_operator:batch_size > self._states[0].size(1); 35, comment:# This batch is larger than the all previous states.; 36, comment:# If so, resize the states.; 37, block; 38, elif_clause; 39, else_clause; 40, comparison_operator:len(self._states) == 1; 41, comment:# GRUs only have a single state. This `unpacks` it from the; 42, comment:# tuple and returns the tensor directly.; 43, block; 44, else_clause; 45, identifier:int; 46, identifier:int; 47, attribute; 48, type; 49, attribute; 50, None; 51, return_statement; 52, identifier:batch_size; 53, call; 54, expression_statement; 55, expression_statement; 56, comment:# state has shape (num_layers, batch_size, hidden_size); 57, for_statement; 58, expression_statement; 59, expression_statement; 60, comparison_operator:batch_size < self._states[0].size(1); 61, comment:# This batch is smaller than the previous one.; 62, block; 63, block; 64, call; 65, integer:1; 66, expression_statement; 67, expression_statement; 68, return_statement; 69, comment:# LSTMs have a state tuple of (state, memory).; 70, block; 71, identifier:torch; 72, identifier:LongTensor; 73, identifier:RnnState; 74, identifier:self; 75, identifier:_states; 76, None; 77, attribute; 78, argument_list; 79, assignment; 80, assignment; 81, identifier:state; 82, attribute; 83, comment:# This _must_ be inside the loop because some; 84, comment:# RNNs have states with different last dimension sizes.; 85, block; 86, assignment; 87, assignment; 88, identifier:batch_size; 89, call; 90, expression_statement; 91, expression_statement; 92, identifier:len; 93, argument_list; 94, assignment; 95, assignment; 96, call; 97, expression_statement; 98, return_statement; 99, subscript; 100, identifier:size; 101, integer:1; 102, identifier:num_states_to_concat; 103, binary_operator:batch_size - self._states[0].size(1); 104, identifier:resized_states; 105, list; 106, identifier:self; 107, identifier:_states; 108, expression_statement; 109, expression_statement; 110, attribute; 111, call; 112, identifier:correctly_shaped_states; 113, attribute; 114, attribute; 115, argument_list; 116, assignment; 117, assignment; 118, attribute; 119, identifier:correctly_shaped_state; 120, subscript; 121, identifier:sorted_state; 122, call; 123, attribute; 124, argument_list; 125, assignment; 126, call; 127, attribute; 128, integer:0; 129, identifier:batch_size; 130, call; 131, assignment; 132, call; 133, identifier:self; 134, identifier:_states; 135, identifier:tuple; 136, argument_list; 137, identifier:self; 138, identifier:_states; 139, subscript; 140, identifier:size; 141, integer:1; 142, identifier:correctly_shaped_states; 143, call; 144, identifier:correctly_shaped_states; 145, attribute; 146, identifier:self; 147, identifier:_states; 148, identifier:correctly_shaped_states; 149, integer:0; 150, attribute; 151, argument_list; 152, subscript; 153, identifier:contiguous; 154, identifier:sorted_states; 155, list_comprehension; 156, identifier:tuple; 157, generator_expression; 158, identifier:self; 159, identifier:_states; 160, attribute; 161, argument_list; 162, identifier:zeros; 163, call; 164, attribute; 165, argument_list; 166, identifier:resized_states; 167, attribute; 168, integer:0; 169, identifier:tuple; 170, generator_expression; 171, identifier:self; 172, identifier:_states; 173, identifier:correctly_shaped_state; 174, identifier:index_select; 175, integer:1; 176, identifier:sorting_indices; 177, identifier:sorted_state; 178, slice; 179, slice; 180, slice; 181, call; 182, for_in_clause; 183, call; 184, for_in_clause; 185, subscript; 186, identifier:size; 187, integer:1; 188, attribute; 189, argument_list; 190, identifier:resized_states; 191, identifier:append; 192, call; 193, identifier:self; 194, identifier:_states; 195, subscript; 196, for_in_clause; 197, identifier:num_valid; 198, attribute; 199, argument_list; 200, identifier:state; 201, identifier:correctly_shaped_states; 202, attribute; 203, argument_list; 204, identifier:state; 205, identifier:sorted_states; 206, attribute; 207, integer:0; 208, identifier:state; 209, identifier:new_zeros; 210, call; 211, identifier:num_states_to_concat; 212, call; 213, attribute; 214, argument_list; 215, identifier:state; 216, slice; 217, slice; 218, slice; 219, identifier:state; 220, attribute; 221, identifier:state; 222, identifier:index_select; 223, integer:1; 224, identifier:sorting_indices; 225, subscript; 226, identifier:contiguous; 227, identifier:self; 228, identifier:_states; 229, attribute; 230, argument_list; 231, attribute; 232, argument_list; 233, identifier:torch; 234, identifier:cat; 235, list; 236, integer:1; 237, identifier:batch_size; 238, identifier:self; 239, identifier:_states; 240, identifier:state; 241, slice; 242, slice; 243, slice; 244, identifier:state; 245, identifier:size; 246, integer:0; 247, identifier:state; 248, identifier:size; 249, integer:2; 250, identifier:state; 251, identifier:zeros; 252, identifier:num_valid
code: def _get_initial_states(self,
                              batch_size: int,
                              num_valid: int,
                              sorting_indices: torch.LongTensor) -> Optional[RnnState]:
          """
          Returns an initial state for use in an RNN. Additionally, this method handles
          the batch size changing across calls by mutating the state to append initial states
          for new elements in the batch. Finally, it also handles sorting the states
          with respect to the sequence lengths of elements in the batch and removing rows
          which are completely padded. Importantly, this `mutates` the state if the
          current batch size is larger than when it was previously called.
  
          Parameters
          ----------
          batch_size : ``int``, required.
              The batch size can change size across calls to stateful RNNs, so we need
              to know if we need to expand or shrink the states before returning them.
              Expanded states will be set to zero.
          num_valid : ``int``, required.
              The batch may contain completely padded sequences which get removed before
              the sequence is passed through the encoder. We also need to clip these off
              of the state too.
          sorting_indices ``torch.LongTensor``, required.
              Pytorch RNNs take sequences sorted by length. When we return the states to be
              used for a given call to ``module.forward``, we need the states to match up to
              the sorted sequences, so before returning them, we sort the states using the
              same indices used to sort the sequences.
  
          Returns
          -------
          This method has a complex return type because it has to deal with the first time it
          is called, when it has no state, and the fact that types of RNN have heterogeneous
          states.
  
          If it is the first time the module has been called, it returns ``None``, regardless
          of the type of the ``Module``.
  
          Otherwise, for LSTMs, it returns a tuple of ``torch.Tensors`` with shape
          ``(num_layers, num_valid, state_size)`` and ``(num_layers, num_valid, memory_size)``
          respectively, or for GRUs, it returns a single ``torch.Tensor`` of shape
          ``(num_layers, num_valid, state_size)``.
          """
          # We don't know the state sizes the first time calling forward,
          # so we let the module define what it's initial hidden state looks like.
          if self._states is None:
              return None
  
          # Otherwise, we have some previous states.
          if batch_size > self._states[0].size(1):
              # This batch is larger than the all previous states.
              # If so, resize the states.
              num_states_to_concat = batch_size - self._states[0].size(1)
              resized_states = []
              # state has shape (num_layers, batch_size, hidden_size)
              for state in self._states:
                  # This _must_ be inside the loop because some
                  # RNNs have states with different last dimension sizes.
                  zeros = state.new_zeros(state.size(0),
                                          num_states_to_concat,
                                          state.size(2))
                  resized_states.append(torch.cat([state, zeros], 1))
              self._states = tuple(resized_states)
              correctly_shaped_states = self._states
  
          elif batch_size < self._states[0].size(1):
              # This batch is smaller than the previous one.
              correctly_shaped_states = tuple(state[:, :batch_size, :] for state in self._states)
          else:
              correctly_shaped_states = self._states
  
          # At this point, our states are of shape (num_layers, batch_size, hidden_size).
          # However, the encoder uses sorted sequences and additionally removes elements
          # of the batch which are fully padded. We need the states to match up to these
          # sorted and filtered sequences, so we do that in the next two blocks before
          # returning the state/s.
          if len(self._states) == 1:
              # GRUs only have a single state. This `unpacks` it from the
              # tuple and returns the tensor directly.
              correctly_shaped_state = correctly_shaped_states[0]
              sorted_state = correctly_shaped_state.index_select(1, sorting_indices)
              return sorted_state[:, :num_valid, :].contiguous()
          else:
              # LSTMs have a state tuple of (state, memory).
              sorted_states = [state.index_select(1, sorting_indices)
                               for state in correctly_shaped_states]
              return tuple(state[:, :num_valid, :].contiguous() for state in sorted_states)