nodes: 0, module; 1, function_definition; 2, function_name:_update_states; 3, parameters; 4, type; 5, block; 6, identifier:self; 7, typed_parameter; 8, typed_parameter; 9, None; 10, expression_statement; 11, comment:# TODO(Mark): seems weird to sort here, but append zeros in the subclasses.; 12, comment:# which way around is best?; 13, expression_statement; 14, if_statement; 15, identifier:final_states; 16, type; 17, identifier:restoration_indices; 18, type; 19, comment:"""
          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.
          """; 20, assignment; 21, comparison_operator:self._states is None; 22, comment:# We don't already have states, so just set the; 23, comment:# ones we receive to be the current state.; 24, block; 25, else_clause; 26, identifier:RnnStateStorage; 27, attribute; 28, identifier:new_unsorted_states; 29, list_comprehension; 30, attribute; 31, None; 32, expression_statement; 33, comment:# Now we've sorted the states back so that they correspond to the original; 34, comment:# indices, we need to figure out what states we need to update, because if we; 35, comment:# didn't use a state for a particular row, we want to preserve its state.; 36, comment:# Thankfully, the rows which are all zero in the state correspond exactly; 37, comment:# to those which aren't used, so we create masks of shape (new_batch_size,),; 38, comment:# denoting which states were used in the RNN computation.; 39, block; 40, identifier:torch; 41, identifier:LongTensor; 42, call; 43, for_in_clause; 44, identifier:self; 45, identifier:_states; 46, assignment; 47, expression_statement; 48, expression_statement; 49, comment:# Masks for the unused states of shape (1, new_batch_size, 1); 50, expression_statement; 51, expression_statement; 52, if_statement; 53, comment:# It looks like there should be another case handled here - when; 54, comment:# the current_state_batch_size < new_state_batch_size. However,; 55, comment:# this never happens, because the states themeselves are mutated; 56, comment:# by appending zeros when calling _get_inital_states, meaning that; 57, comment:# the new states are either of equal size, or smaller, in the case; 58, comment:# that there are some unused elements (zero-length) for the RNN computation.; 59, expression_statement; 60, attribute; 61, argument_list; 62, identifier:state; 63, identifier:final_states; 64, attribute; 65, call; 66, assignment; 67, assignment; 68, assignment; 69, assignment; 70, comparison_operator:current_state_batch_size > new_state_batch_size; 71, comment:# The new state is smaller than the old one,; 72, comment:# so just update the indices which we used.; 73, block; 74, else_clause; 75, assignment; 76, identifier:state; 77, identifier:index_select; 78, integer:1; 79, identifier:restoration_indices; 80, identifier:self; 81, identifier:_states; 82, identifier:tuple; 83, generator_expression; 84, identifier:current_state_batch_size; 85, call; 86, identifier:new_state_batch_size; 87, call; 88, identifier:used_new_rows_mask; 89, list_comprehension; 90, identifier:new_states; 91, list; 92, identifier:current_state_batch_size; 93, identifier:new_state_batch_size; 94, for_statement; 95, comment:# The states are the same size, so we just have to; 96, comment:# deal with the possibility that some rows weren't used.; 97, block; 98, attribute; 99, call; 100, attribute; 101, for_in_clause; 102, attribute; 103, argument_list; 104, attribute; 105, argument_list; 106, call; 107, for_in_clause; 108, pattern_list; 109, call; 110, comment:# zero out all rows in the previous state; 111, comment:# which _were_ used in the current state.; 112, block; 113, expression_statement; 114, for_statement; 115, identifier:self; 116, identifier:_states; 117, identifier:tuple; 118, argument_list; 119, identifier:state; 120, identifier:data; 121, identifier:state; 122, identifier:new_unsorted_states; 123, subscript; 124, identifier:size; 125, integer:1; 126, subscript; 127, identifier:size; 128, integer:1; 129, attribute; 130, argument_list; 131, identifier:state; 132, identifier:new_unsorted_states; 133, identifier:old_state; 134, identifier:new_state; 135, identifier:used_mask; 136, identifier:zip; 137, argument_list; 138, expression_statement; 139, comment:# The old state is larger, so update the relevant parts of it.; 140, expression_statement; 141, expression_statement; 142, assignment; 143, pattern_list; 144, call; 145, comment:# zero out all rows which _were_ used in the current state.; 146, block; 147, identifier:new_states; 148, attribute; 149, integer:0; 150, identifier:final_states; 151, integer:0; 152, call; 153, identifier:view; 154, integer:1; 155, identifier:new_state_batch_size; 156, integer:1; 157, attribute; 158, identifier:new_unsorted_states; 159, identifier:used_new_rows_mask; 160, assignment; 161, assignment; 162, call; 163, identifier:new_states; 164, list; 165, identifier:old_state; 166, identifier:new_state; 167, identifier:used_mask; 168, identifier:zip; 169, argument_list; 170, expression_statement; 171, comment:# The old state is larger, so update the relevant parts of it.; 172, expression_statement; 173, expression_statement; 174, identifier:self; 175, identifier:_states; 176, attribute; 177, argument_list; 178, identifier:self; 179, identifier:_states; 180, identifier:masked_old_state; 181, binary_operator:old_state[:, :new_state_batch_size, :] * (1 - used_mask); 182, subscript; 183, binary_operator:new_state + masked_old_state; 184, attribute; 185, argument_list; 186, attribute; 187, identifier:new_unsorted_states; 188, identifier:used_new_rows_mask; 189, assignment; 190, augmented_assignment; 191, call; 192, parenthesized_expression; 193, identifier:float; 194, subscript; 195, parenthesized_expression; 196, identifier:old_state; 197, slice; 198, slice; 199, slice; 200, identifier:new_state; 201, identifier:masked_old_state; 202, identifier:new_states; 203, identifier:append; 204, call; 205, identifier:self; 206, identifier:_states; 207, identifier:masked_old_state; 208, binary_operator:old_state * (1 - used_mask); 209, identifier:new_state; 210, identifier:masked_old_state; 211, attribute; 212, argument_list; 213, comparison_operator:state[0, :, :].sum(-1)
                                     != 0.0; 214, identifier:old_state; 215, slice; 216, slice; 217, slice; 218, binary_operator:1 - used_mask; 219, identifier:new_state_batch_size; 220, attribute; 221, argument_list; 222, identifier:old_state; 223, parenthesized_expression; 224, identifier:new_states; 225, identifier:append; 226, call; 227, call; 228, float:0.0; 229, identifier:new_state_batch_size; 230, integer:1; 231, identifier:used_mask; 232, identifier:old_state; 233, identifier:detach; 234, binary_operator:1 - used_mask; 235, attribute; 236, argument_list; 237, attribute; 238, argument_list; 239, integer:1; 240, identifier:used_mask; 241, identifier:new_state; 242, identifier:detach; 243, subscript; 244, identifier:sum; 245, unary_operator; 246, identifier:state; 247, integer:0; 248, slice; 249, slice; 250, integer:1
code: def _update_states(self,
                         final_states: RnnStateStorage,
                         restoration_indices: torch.LongTensor) -> None:
          """
          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 zip(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 zip(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)