code: def test_update_states(self):
      assert self.encoder_base._states is None
      initial_states = torch.randn([1, 5, 7]), torch.randn([1, 5, 7])
  
      index_selected_initial_states = (
          initial_states[0].index_select(1, self.restoration_indices),
          initial_states[1].index_select(1, self.restoration_indices),
      )
  
      self.encoder_base._update_states(
          initial_states, self.restoration_indices
      )
      # State was None, so the updated state should just be the sorted given state.
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0].data.numpy(),
          index_selected_initial_states[0].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1].data.numpy(),
          index_selected_initial_states[1].data.numpy(),
      )
  
      new_states = torch.randn([1, 5, 7]), torch.randn([1, 5, 7])
      # tensor has 2 completely masked rows, so the last 2 rows of the _sorted_ states
      # will be completely zero, having been appended after calling the respective encoder.
      new_states[0][:, -2:, :] = 0
      new_states[1][:, -2:, :] = 0
  
      index_selected_new_states = (
          new_states[0].index_select(1, self.restoration_indices),
          new_states[1].index_select(1, self.restoration_indices),
      )
  
      self.encoder_base._update_states(
          new_states, self.restoration_indices
      )
      # Check that the update _preserved_ the state for the rows which were
      # completely masked (2 and 4):
      for index in [2, 4]:
          numpy.testing.assert_array_equal(
              self.encoder_base._states[0][:, index, :].data.numpy(),
              index_selected_initial_states[0][
                  :, index, :
              ].data.numpy(),
          )
          numpy.testing.assert_array_equal(
              self.encoder_base._states[1][:, index, :].data.numpy(),
              index_selected_initial_states[1][
                  :, index, :
              ].data.numpy(),
          )
      # Now the states which were updated:
      for index in [0, 1, 3]:
          numpy.testing.assert_array_equal(
              self.encoder_base._states[0][:, index, :].data.numpy(),
              index_selected_new_states[0][:, index, :].data.numpy(),
          )
          numpy.testing.assert_array_equal(
              self.encoder_base._states[1][:, index, :].data.numpy(),
              index_selected_new_states[1][:, index, :].data.numpy(),
          )
  
      # Now test the case that the new state is smaller:
      small_new_states = torch.randn([1, 3, 7]), torch.randn([1, 3, 7])
      # pretend the 2nd sequence in the batch was fully masked.
      small_restoration_indices = torch.LongTensor([2, 0, 1])
      small_new_states[0][:, 0, :] = 0
      small_new_states[1][:, 0, :] = 0
  
      index_selected_small_states = (
          small_new_states[0].index_select(
              1, small_restoration_indices
          ),
          small_new_states[1].index_select(
              1, small_restoration_indices
          ),
      )
      self.encoder_base._update_states(
          small_new_states, small_restoration_indices
      )
  
      # Check the index for the row we didn't update is the same as the previous step:
      for index in [1, 3]:
          numpy.testing.assert_array_equal(
              self.encoder_base._states[0][:, index, :].data.numpy(),
              index_selected_new_states[0][:, index, :].data.numpy(),
          )
          numpy.testing.assert_array_equal(
              self.encoder_base._states[1][:, index, :].data.numpy(),
              index_selected_new_states[1][:, index, :].data.numpy(),
          )
      # Indices we did update:
      for index in [0, 2]:
          numpy.testing.assert_array_equal(
              self.encoder_base._states[0][:, index, :].data.numpy(),
              index_selected_small_states[0][:, index, :].data.numpy(),
          )
          numpy.testing.assert_array_equal(
              self.encoder_base._states[1][:, index, :].data.numpy(),
              index_selected_small_states[1][:, index, :].data.numpy(),
          )
  
      # We didn't update index 4 in the previous step either, so it should be equal to the
      # 4th index of initial states.
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0][:, 4, :].data.numpy(),
          index_selected_initial_states[0][:, 4, :].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1][:, 4, :].data.numpy(),
          index_selected_initial_states[1][:, 4, :].data.numpy(),
      )
  
file_path: magnitude/pymagnitude/third_party/allennlp/tests/modules/encoder_base_test.py
imports: from __future__ import absolute_import
  import numpy
  import torch
  from torch.nn import LSTM, RNN
  from allennlp.modules.encoder_base import _EncoderBase
  from allennlp.common.testing import AllenNlpTestCase
  from allennlp.nn.util import sort_batch_by_length, get_lengths_from_binary_sequence_mask
label_desc: Use image node_id 4 for calling the TestEncoderBase obj's underlying member method code with example usage: obj.test_update_states() without return types