code: def _test_n_step_forward(self, gpu):
      in_size = 2
      out0_size = 3
      out1_size = 4
      out2_size = 1
  
      par = StatelessRecurrentBranched(
          L.NStepLSTM(1, in_size, out0_size, 0),
          StatelessRecurrentSequential(
              L.NStepRNNReLU(1, in_size, out1_size, 0),
          ),
          StatelessRecurrentSequential(
              L.Linear(in_size, out2_size),
          ),
      )
  
      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=(1, in_size)).astype(
              np.float32
          ),
          xp.random.uniform(-1, 1, size=(3, in_size)).astype(
              np.float32
          ),
      ]
  
      # Concatenated output should be a tuple of three variables.
      concat_out, concat_rs = par.n_step_forward(
          seqs_x, None, output_mode="concat"
      )
      self.assertIsInstance(concat_out, tuple)
      self.assertEqual(len(concat_out), len(par))
      self.assertEqual(concat_out[0].shape, (4, out0_size))
      self.assertEqual(concat_out[1].shape, (4, out1_size))
      self.assertEqual(concat_out[2].shape, (4, out2_size))
  
      self.assertIsInstance(concat_rs, tuple)
      self.assertEqual(len(concat_rs), len(par))
      self.assertIsInstance(concat_rs[0], tuple)
      # NStepLSTM
      self.assertEqual(len(concat_rs[0]), 2)
      self.assertEqual(
          concat_rs[0][0].shape, (1, len(seqs_x), out0_size)
      )
      self.assertEqual(
          concat_rs[0][1].shape, (1, len(seqs_x), out0_size)
      )
      # StatelessRecurrentSequential(NStepRNNReLU)
      self.assertEqual(len(concat_rs[1]), 1)
      self.assertEqual(
          concat_rs[1][0].shape, (1, len(seqs_x), out1_size)
      )
      # StatelessRecurrentSequential(Linear)
      self.assertEqual(len(concat_rs[2]), 0)
  
      # Split output should be a list of two tuples, each of which is of
      # three variables.
      split_out, split_rs = par.n_step_forward(
          seqs_x, None, output_mode="split"
      )
      self.assertIsInstance(split_out, list)
      self.assertEqual(len(split_out), len(seqs_x))
      self.assertEqual(len(split_out[0]), len(par))
      self.assertEqual(len(split_out[1]), len(par))
      self.assertEqual(
          split_out[0][0].shape,
          (
              1,
              out0_size,
          ),
      )
      self.assertEqual(
          split_out[0][1].shape,
          (
              1,
              out1_size,
          ),
      )
      self.assertEqual(
          split_out[0][2].shape,
          (
              1,
              out2_size,
          ),
      )
      self.assertEqual(
          split_out[1][0].shape,
          (
              3,
              out0_size,
          ),
      )
      self.assertEqual(
          split_out[1][1].shape,
          (
              3,
              out1_size,
          ),
      )
      self.assertEqual(
          split_out[1][2].shape,
          (
              3,
              out2_size,
          ),
      )
  
      # Check if output_mode='concat' and output_mode='split' are consistent
      xp.testing.assert_allclose(
          F.concat(
              [F.concat(seq_out, axis=1) for seq_out in split_out],
              axis=0,
          ).array,
          F.concat(concat_out, axis=1).array,
      )
  
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 1 for calling the TestStatelessRecurrentBranched obj's underlying member method code with example usage: obj._test_n_step_forward(gpu) without return types