code: def test_get_initial_states(self):
      # First time we call it, there should be no state, so we should return None.
      assert (
          self.encoder_base._get_initial_states(
              self.batch_size, self.num_valid, self.sorting_indices
          )
          is None
      )
  
      # First test the case that the previous state is _smaller_ than the current state input.
      initial_states = (torch.randn([1, 3, 7]), torch.randn([1, 3, 7]))
      self.encoder_base._states = initial_states
      # sorting indices are: [0, 1, 3, 2, 4]
      returned_states = self.encoder_base._get_initial_states(
          self.batch_size, self.num_valid, self.sorting_indices
      )
  
      correct_expanded_states = [
          torch.cat([state, torch.zeros([1, 2, 7])], 1)
          for state in initial_states
      ]
      # State should have been expanded with zeros to have shape (1, batch_size, hidden_size).
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0].data.numpy(),
          correct_expanded_states[0].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1].data.numpy(),
          correct_expanded_states[1].data.numpy(),
      )
  
      # The returned states should be of shape (1, num_valid, hidden_size) and
      # they also should have been sorted with respect to the indices.
      # sorting indices are: [0, 1, 3, 2, 4]
  
      correct_returned_states = [
          state.index_select(1, self.sorting_indices)[
              :, : self.num_valid, :
          ]
          for state in correct_expanded_states
      ]
  
      numpy.testing.assert_array_equal(
          returned_states[0].data.numpy(),
          correct_returned_states[0].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          returned_states[1].data.numpy(),
          correct_returned_states[1].data.numpy(),
      )
  
      # Now test the case that the previous state is larger:
      original_states = (
          torch.randn([1, 10, 7]),
          torch.randn([1, 10, 7]),
      )
      self.encoder_base._states = original_states
      # sorting indices are: [0, 1, 3, 2, 4]
      returned_states = self.encoder_base._get_initial_states(
          self.batch_size, self.num_valid, self.sorting_indices
      )
      # State should not have changed, as they were larger
      # than the batch size of the requested states.
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0].data.numpy(),
          original_states[0].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1].data.numpy(),
          original_states[1].data.numpy(),
      )
  
      # The returned states should be of shape (1, num_valid, hidden_size) and they
      # also should have been sorted with respect to the indices.
      correct_returned_state = [
          x.index_select(1, self.sorting_indices)[
              :, : self.num_valid, :
          ]
          for x in original_states
      ]
      numpy.testing.assert_array_equal(
          returned_states[0].data.numpy(),
          correct_returned_state[0].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          returned_states[1].data.numpy(),
          correct_returned_state[1].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 3 for calling the TestEncoderBase obj's underlying member method code with example usage: obj.test_get_initial_states() without return types