class_docstr: This abstract class serves as a base for the 3 ``Encoder`` abstractions in AllenNLP.
  - :class:`~allennlp.modules.seq2seq_encoders.Seq2SeqEncoders`
  - :class:`~allennlp.modules.seq2vec_encoders.Seq2VecEncoders`
  
  Additionally, this class provides functionality for sorting sequences by length
  so they can be consumed by Pytorch RNN classes, which require their inputs to be
  sorted by length. Finally, it also provides optional statefulness to all of it's
  subclasses by allowing the caching and retrieving of the hidden states of RNNs.
docstr: After the RNN has run forward, the states need to be updated.
  This method just sets the state to the updated new state, performing
  several pieces of book-keeping along the way - namely, unsorting the
  states and ensuring that the states of completely padded sequences are
  not updated. Finally, it also detaches the state variable from the
  computational graph, such that the graph can be garbage collected after
  each batch iteration.
  
  Parameters
  ----------
  final_states : ``RnnStateStorage``, required.
      The hidden states returned as output from the RNN.
  restoration_indices : ``torch.LongTensor``, required.
      The indices that invert the sorting used in ``sort_and_run_forward``
      to order the states with respect to the lengths of the sequences in
      the batch.
docstr_tok: ["After","the","RNN","has","run","forward",",","the","states","need","to","be","updated",".","This","method","just","sets","the","state","to","the","updated","new","state",",","performing","several","pieces","of","book-keeping","along","the","way","-","namely",",","unsorting","the","states","and","ensuring","that","the","states","of","completely","padded","sequences","are","not","updated",".","Finally",",","it","also","detaches","the","state","variable","from","the","computational","graph",",","such","that","the","graph","can","be","garbage","collected","after","each","batch","iteration",".","Parameters","--","--","--","--","--","final_states",":","``","RnnStateStorage","``",",","required",".","The","hidden","states","returned","as","output","from","the","RNN",".","restoration_indices",":","``","torch.LongTensor","``",",","required",".","The","indices","that","invert","the","sorting","used","in","``","sort_and_run_forward","``","to","order","the","states","with","respect","to","the","lengths","of","the","sequences","in","the","batch","."]
code: def _update_states(self, final_states, restoration_indices):
      """
      After the RNN has run forward, the states need to be updated.
      This method just sets the state to the updated new state, performing
      several pieces of book-keeping along the way - namely, unsorting the
      states and ensuring that the states of completely padded sequences are
      not updated. Finally, it also detaches the state variable from the
      computational graph, such that the graph can be garbage collected after
      each batch iteration.
  
      Parameters
      ----------
      final_states : ``RnnStateStorage``, required.
          The hidden states returned as output from the RNN.
      restoration_indices : ``torch.LongTensor``, required.
          The indices that invert the sorting used in ``sort_and_run_forward``
          to order the states with respect to the lengths of the sequences in
          the batch.
      """
      # TODO(Mark): seems weird to sort here, but append zeros in the subclasses.
      # which way around is best?
      new_unsorted_states = [
          state.index_select(1, restoration_indices)
          for state in final_states
      ]
  
      if self._states is None:
          # We don't already have states, so just set the
          # ones we receive to be the current state.
          self._states = tuple(
              state.data for state in new_unsorted_states
          )
      else:
          # Now we've sorted the states back so that they correspond to the original
          # indices, we need to figure out what states we need to update, because if we
          # didn't use a state for a particular row, we want to preserve its state.
          # Thankfully, the rows which are all zero in the state correspond exactly
          # to those which aren't used, so we create masks of shape (new_batch_size,),
          # denoting which states were used in the RNN computation.
          current_state_batch_size = self._states[0].size(1)
          new_state_batch_size = final_states[0].size(1)
          # Masks for the unused states of shape (1, new_batch_size, 1)
          used_new_rows_mask = [
              (state[0, :, :].sum(-1) != 0.0)
              .float()
              .view(1, new_state_batch_size, 1)
              for state in new_unsorted_states
          ]
          new_states = []
          if current_state_batch_size > new_state_batch_size:
              # The new state is smaller than the old one,
              # so just update the indices which we used.
              for old_state, new_state, used_mask in izip(
                  self._states, new_unsorted_states, used_new_rows_mask
              ):
                  # zero out all rows in the previous state
                  # which _were_ used in the current state.
                  masked_old_state = old_state[
                      :, :new_state_batch_size, :
                  ] * (1 - used_mask)
                  # The old state is larger, so update the relevant parts of it.
                  old_state[:, :new_state_batch_size, :] = (
                      new_state + masked_old_state
                  )
                  new_states.append(old_state.detach())
          else:
              # The states are the same size, so we just have to
              # deal with the possibility that some rows weren't used.
              new_states = []
              for old_state, new_state, used_mask in izip(
                  self._states, new_unsorted_states, used_new_rows_mask
              ):
                  # zero out all rows which _were_ used in the current state.
                  masked_old_state = old_state * (1 - used_mask)
                  # The old state is larger, so update the relevant parts of it.
                  new_state += masked_old_state
                  new_states.append(new_state.detach())
  
          # It looks like there should be another case handled here - when
          # the current_state_batch_size < new_state_batch_size. However,
          # this never happens, because the states themeselves are mutated
          # by appending zeros when calling _get_inital_states, meaning that
          # the new states are either of equal size, or smaller, in the case
          # that there are some unused elements (zero-length) for the RNN computation.
          self._states = tuple(new_states)
  
file_path: magnitude/pymagnitude/third_party/allennlp/modules/encoder_base.py
imports: from __future__ import absolute_import
  import torch
  from torch.nn.utils.rnn import pack_padded_sequence, PackedSequence
  from allennlp.nn.util import get_lengths_from_binary_sequence_mask, sort_batch_by_length
label_desc: Use image node_id 4 for calling the _EncoderBase obj's underlying member method code with example usage: obj._update_states(final_states, restoration_indices) without return types