code: def test_reset_states(self):
      # Initialize the encoder states.
      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
      )
  
      # Check that only some of the states are reset when a mask is provided.
      mask = torch.tensor([True, True, False, False, False])
      self.encoder_base.reset_states(mask)
      # First two states should be zeros
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0][:, :2, :].data.numpy(),
          torch.zeros_like(initial_states[0])[:, :2, :].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1][:, :2, :].data.numpy(),
          torch.zeros_like(initial_states[1])[:, :2, :].data.numpy(),
      )
      # Remaining states should be the same
      numpy.testing.assert_array_equal(
          self.encoder_base._states[0][:, 2:, :].data.numpy(),
          index_selected_initial_states[0][:, 2:, :].data.numpy(),
      )
      numpy.testing.assert_array_equal(
          self.encoder_base._states[1][:, 2:, :].data.numpy(),
          index_selected_initial_states[1][:, 2:, :].data.numpy(),
      )
  
      # Check that error is raised if mask has wrong batch size.
      bad_mask = torch.tensor([True, True, False])
      with pytest.raises(ValueError):
          self.encoder_base.reset_states(bad_mask)
  
      # Check that states are reset to None if no mask is provided.
      self.encoder_base.reset_states()
      assert self.encoder_base._states is None
  
file_path: allennlp/tests/modules/encoder_base_test.py
imports: import numpy
  import pytest
  import torch
  from torch.nn import LSTM, RNN
  from allennlp.modules.encoder_base import _EncoderBase
  from allennlp.common.testing import AllenNlpTestCase, requires_gpu
  from allennlp.nn.util import sort_batch_by_length, get_lengths_from_binary_sequence_mask
label_desc: Use image node_id 5 for calling the TestEncoderBase obj's underlying member method code with example usage: obj.test_reset_states() without return types