code: def _test_mask_recurrent_state_at(self, gpu):
      in_size = 2
      out0_size = 2
      out1_size = 3
      par = StatelessRecurrentBranched(
          L.NStepGRU(1, in_size, out0_size, 0),
          StatelessRecurrentSequential(
              L.NStepLSTM(1, in_size, out1_size, 0),
          ),
      )
      if gpu >= 0:
          chainer.cuda.get_device_from_id(gpu).use()
          par.to_gpu()
      xp = par.xp
      seqs_x = [
          xp.random.uniform(-1, 1, size=(2, in_size)).astype(
              np.float32
          ),
          xp.random.uniform(-1, 1, size=(2, in_size)).astype(
              np.float32
          ),
      ]
      transposed_x = F.transpose_sequence(seqs_x)
  
      nstep_out, nstep_rs = par.n_step_forward(
          seqs_x, None, output_mode="concat"
      )
  
      # Check if n_step_forward and forward twice results are same
      def no_mask_forward_twice():
          _, rs = par(transposed_x[0], None)
          return par(transposed_x[1], rs)
  
      nomask_out, nomask_rs = no_mask_forward_twice()
      # GRU
      xp.testing.assert_allclose(
          nstep_out[0].array[[1, 3]],
          nomask_out[0].array,
      )
      # LSTM
      xp.testing.assert_allclose(
          nstep_out[1].array[[1, 3]],
          nomask_out[1].array,
      )
      xp.testing.assert_allclose(nstep_rs[0].array, nomask_rs[0].array)
      self.assertIsInstance(nomask_rs[1], tuple)
      self.assertEqual(len(nomask_rs[1]), 1)
      self.assertEqual(len(nomask_rs[1][0]), 2)
      xp.testing.assert_allclose(
          nstep_rs[1][0][0].array, nomask_rs[1][0][0].array
      )
      xp.testing.assert_allclose(
          nstep_rs[1][0][1].array, nomask_rs[1][0][1].array
      )
  
      # 1st-only mask forward twice: only 2nd should be the same
      def mask0_forward_twice():
          _, rs = par(transposed_x[0], None)
          rs = par.mask_recurrent_state_at(rs, 0)
          return par(transposed_x[1], rs)
  
      mask0_out, mask0_rs = mask0_forward_twice()
      # GRU
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[0].array[1],
              mask0_out[0].array[0],
          )
      xp.testing.assert_allclose(
          nstep_out[0].array[3],
          mask0_out[0].array[1],
      )
      # LSTM
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[1].array[1],
              mask0_out[1].array[0],
          )
      xp.testing.assert_allclose(
          nstep_out[1].array[3],
          mask0_out[1].array[1],
      )
  
      # 2nd-only mask forward twice: only 1st should be the same
      def mask1_forward_twice():
          _, rs = par(transposed_x[0], None)
          rs = par.mask_recurrent_state_at(rs, 1)
          return par(transposed_x[1], rs)
  
      mask1_out, mask1_rs = mask1_forward_twice()
      # GRU
      xp.testing.assert_allclose(
          nstep_out[0].array[1],
          mask1_out[0].array[0],
      )
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[0].array[3],
              mask1_out[0].array[1],
          )
      # LSTM
      xp.testing.assert_allclose(
          nstep_out[1].array[1],
          mask1_out[1].array[0],
      )
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[1].array[3],
              mask1_out[1].array[1],
          )
  
      # both 1st and 2nd mask forward twice: both should be different
      def mask01_forward_twice():
          _, rs = par(transposed_x[0], None)
          rs = par.mask_recurrent_state_at(rs, [0, 1])
          return par(transposed_x[1], rs)
  
      mask01_out, mask01_rs = mask01_forward_twice()
      # GRU
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[0].array[1],
              mask01_out[0].array[0],
          )
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[0].array[3],
              mask01_out[0].array[1],
          )
      # LSTM
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[1].array[1],
              mask01_out[1].array[0],
          )
      with self.assertRaises(AssertionError):
          xp.testing.assert_allclose(
              nstep_out[1].array[3],
              mask01_out[1].array[1],
          )
  
      # get and concat recurrent states and resume forward
      def get_and_concat_rs_forward():
          _, rs = par(transposed_x[0], None)
          rs0 = par.get_recurrent_state_at(rs, 0, unwrap_variable=True)
          rs1 = par.get_recurrent_state_at(rs, 1, unwrap_variable=True)
          concat_rs = par.concatenate_recurrent_states([rs0, rs1])
          return par(transposed_x[1], concat_rs)
  
      getcon_out, getcon_rs = get_and_concat_rs_forward()
      # GRU
      xp.testing.assert_allclose(
          nstep_out[0].array[1],
          getcon_out[0].array[0],
      )
      xp.testing.assert_allclose(
          nstep_out[0].array[3],
          getcon_out[0].array[1],
      )
      # LSTM
      xp.testing.assert_allclose(
          nstep_out[1].array[1],
          getcon_out[1].array[0],
      )
      xp.testing.assert_allclose(
          nstep_out[1].array[3],
          getcon_out[1].array[1],
      )
  
file_path: chainerrl/tests/links_tests/test_stateless_recurrent_branched.py
imports: import unittest
  import chainer
  from chainer import functions
  from chainer import links
  from chainer import testing
  import numpy
  from chainerrl.links import StatelessRecurrentBranched
  from chainerrl.links import StatelessRecurrentSequential
label_desc: Use image node_id 4 for calling the TestStatelessRecurrentBranched obj's underlying member method code with example usage: obj._test_mask_recurrent_state_at(gpu) and returns: par, par, par, par, par