code: import unittest
  
  import mock
  import numpy
  
  import chainer
  from chainer import cuda
  from chainer import functions
  from chainer import gradient_check
  from chainer import testing
  from chainer.testing import attr
  
  
  def sigmoid(x):
      return numpy.tanh(x * 0.5) * 0.5 + 0.5
  
  
  def _split(inputs, pos):
      return inputs[:pos], inputs[pos:]
  
  
  @testing.parameterize(*testing.product({
      'use_cudnn': [True, False],
  }))
  class TestNStepLSTM(unittest.TestCase):
  
      batches = [3, 2, 1]
      length = len(batches)
      in_size = 3
      out_size = 2
      n_layers = 2
      dropout = 0.0
  
      def setUp(self):
          self.xs = [numpy.random.uniform(-1, 1, (b, self.in_size)).astype('f')
                     for b in self.batches]
          h_shape = (self.n_layers, self.batches[0], self.out_size)
          self.cx = numpy.random.uniform(-1, 1, h_shape).astype(numpy.float32)
          self.hx = numpy.random.uniform(-1, 1, h_shape).astype(numpy.float32)
  
          self.ws = []
          self.bs = []
          for i in range(self.n_layers):
              weights = []
              biases = []
              for j in range(8):
                  if i == 0 and j < 4:
                      w_in = self.in_size
                  else:
                      w_in = self.out_size
  
                  weights.append(numpy.random.uniform(
                      -1, 1, (self.out_size, w_in)).astype('f'))
                  biases.append(numpy.random.uniform(
                      -1, 1, (self.out_size,)).astype('f'))
              self.ws.append(weights)
              self.bs.append(biases)
  
          self.dys = [numpy.random.uniform(-1, 1, (b, self.out_size)).astype('f')
                      for b in self.batches]
          self.dcy = numpy.random.uniform(-1, 1, h_shape).astype(numpy.float32)
          self.dhy = numpy.random.uniform(-1, 1, h_shape).astype(numpy.float32)
  
      def check_forward(
              self, h_data, c_data, xs_data, ws_data, bs_data, volatile):
          h = chainer.Variable(h_data, volatile=volatile)
          c = chainer.Variable(c_data, volatile=volatile)
          xs = [chainer.Variable(x, volatile=volatile) for x in xs_data]
          ws = [[chainer.Variable(w, volatile=volatile) for w in ws]
                for ws in ws_data]
          bs = [[chainer.Variable(b, volatile=volatile) for b in bs]
                for bs in bs_data]
          hy, cy, ys = functions.n_step_lstm(
              self.n_layers, self.dropout, h, c, ws, bs, xs,
              use_cudnn=self.use_cudnn)
  
          e_hy = self.hx.copy()
          e_cy = self.cx.copy()
          for ind in range(self.length):
              x = self.xs[ind]
              batch = x.shape[0]
              for layer in range(self.n_layers):
                  w = self.ws[layer]
                  b = self.bs[layer]
                  h_prev = e_hy[layer, :batch]
                  c_prev = e_cy[layer, :batch]
                  i = sigmoid(x.dot(w[0].T) + h_prev.dot(w[4].T) + b[0] + b[4])
                  f = sigmoid(x.dot(w[1].T) + h_prev.dot(w[5].T) + b[1] + b[5])
                  c_bar = numpy.tanh(
                      x.dot(w[2].T) + h_prev.dot(w[6].T) + b[2] + b[6])
                  o = sigmoid(x.dot(w[3].T) + h_prev.dot(w[7].T) + b[3] + b[7])
                  e_c = (f * c_prev + i * c_bar)
                  e_h = o * numpy.tanh(e_c)
                  e_hy[layer, :batch] = e_h
                  e_cy[layer, :batch] = e_c
  
                  x = e_h
  
              testing.assert_allclose(
                  ys[ind].data, x, rtol=1e-4, atol=1e-4)
  
          testing.assert_allclose(hy.data, e_hy, rtol=1e-4, atol=1e-4)
          testing.assert_allclose(cy.data, e_cy, rtol=1e-4, atol=1e-4)
  
      def test_forward_cpu(self):
          self.check_forward(self.hx, self.cx, self.xs, self.ws, self.bs, False)
  
      def test_forward_cpu_volatile(self):
          self.check_forward(self.hx, self.cx, self.xs, self.ws, self.bs, True)
  
      @attr.gpu
      def test_forward_gpu(self):
          self.check_forward(cuda.to_gpu(self.hx),
                             cuda.to_gpu(self.cx),
                             [cuda.to_gpu(x) for x in self.xs],
                             [[cuda.to_gpu(w) for w in ws] for ws in self.ws],
                             [[cuda.to_gpu(b) for b in bs] for bs in self.bs],
                             False)
  
      @attr.gpu
      def test_forward_gpu_volatile(self):
          self.check_forward(cuda.to_gpu(self.hx),
                             cuda.to_gpu(self.cx),
                             [cuda.to_gpu(x) for x in self.xs],
                             [[cuda.to_gpu(w) for w in ws] for ws in self.ws],
                             [[cuda.to_gpu(b) for b in bs] for bs in self.bs],
                             True)
  
      def check_backward(self, h_data, c_data, xs_data, ws_data, bs_data,
                         dhy_data, dcy_data, dys_data):
          args = tuple([h_data, c_data] + sum(ws_data, []) + sum(bs_data, []) +
                       xs_data)
          grads = tuple([dhy_data, dcy_data] + dys_data)
  
          def f(*inputs):
              (hx, cx), inputs = _split(inputs, 2)
              ws = []
              for i in range(self.n_layers):
                  weights, inputs = _split(inputs, 8)
                  ws.append(weights)
              bs = []
              for i in range(self.n_layers):
                  biases, inputs = _split(inputs, 8)
                  bs.append(biases)
              xs = inputs
              hy, cy, ys = functions.n_step_lstm(
                  self.n_layers, self.dropout, hx, cx, ws, bs, xs)
              return (hy, cy) + ys
  
          gradient_check.check_backward(
              f, args, grads, eps=1e-2, rtol=1e-3, atol=1e-3)
  
      def test_backward_cpu(self):
          self.check_backward(self.hx, self.cx, self.xs, self.ws, self.bs,
                              self.dhy, self.dcy, self.dys)
  
      @attr.gpu
      def test_backward_gpu(self):
          self.check_backward(cuda.to_gpu(self.hx),
                              cuda.to_gpu(self.cx),
                              [cuda.to_gpu(x) for x in self.xs],
                              [[cuda.to_gpu(w) for w in ws] for ws in self.ws],
                              [[cuda.to_gpu(b) for b in bs] for bs in self.bs],
                              cuda.to_gpu(self.dhy),
                              cuda.to_gpu(self.dcy),
                              [cuda.to_gpu(dy) for dy in self.dys])
  
  
  @testing.parameterize(*testing.product({
      'use_cudnn': [True, False],
  }))
  @attr.cudnn
  class TestNStepLSTMCudnnCall(unittest.TestCase):
  
      batches = [4, 3, 2, 1]
      length = len(batches)
      in_size = 3
      out_size = 4
      n_layers = 2
      dropout = 0.0
  
      def setUp(self):
          self.xs = [cuda.cupy.random.uniform(
              -1, 1, (b, self.in_size)).astype('f')
              for b in self.batches]
          h_shape = (self.n_layers, self.batches[0], self.out_size)
          self.cx = cuda.cupy.random.uniform(-1, 1, h_shape).astype('f')
          self.hx = cuda.cupy.random.uniform(-1, 1, h_shape).astype('f')
  
          self.ws = []
          self.bs = []
          for i in range(self.n_layers):
              weights = []
              biases = []
              for j in range(8):
                  if i == 0 and j < 4:
                      w_in = self.in_size
                  else:
                      w_in = self.out_size
  
                  weights.append(cuda.cupy.random.uniform(
                      -1, 1, (self.out_size, w_in)).astype('f'))
                  biases.append(cuda.cupy.random.uniform(
                      -1, 1, (self.out_size,)).astype('f'))
  
              self.ws.append(weights)
              self.bs.append(biases)
  
          self.dys = [cuda.cupy.random.uniform(
              -1, 1, (b, self.out_size)).astype('f')
              for b in self.batches]
          self.dcy = cuda.cupy.random.uniform(-1, 1, h_shape).astype('f')
          self.dhy = cuda.cupy.random.uniform(-1, 1, h_shape).astype('f')
          self.expect = self.use_cudnn and (
              cuda.cudnn.cudnn.getVersion() >= 5000)
  
      def forward(self, train):
          volatile = not train
          h = chainer.Variable(self.hx, volatile=volatile)
          c = chainer.Variable(self.cx, volatile=volatile)
          xs = [chainer.Variable(x, volatile=volatile) for x in self.xs]
          ws = [[chainer.Variable(w, volatile=volatile) for w in ws]
                for ws in self.ws]
          bs = [[chainer.Variable(b, volatile=volatile) for b in bs]
                for bs in self.bs]
          with chainer.using_config('train', train):
              return functions.n_step_lstm(
                  self.n_layers, self.dropout, h, c, ws, bs, xs,
                  use_cudnn=self.use_cudnn)
  
      def test_call_cudnn_forward_training(self):
          with mock.patch('cupy.cuda.cudnn.RNNForwardTraining') as func:
              self.forward(True)
              self.assertEqual(func.called, self.expect)
  
      def test_call_cudnn_forward_inference(self):
          with mock.patch('cupy.cuda.cudnn.RNNForwardInference') as func:
              self.forward(False)
              self.assertEqual(func.called, self.expect)
  
      def test_call_cudnn_backward(self):
          hy, cy, ys = self.forward(True)
          hy.grad = self.dhy
          with mock.patch('cupy.cuda.cudnn.RNNBackwardWeights') as func:
              hy.backward()
              self.assertEqual(func.called, self.expect)
  
  
  testing.run_module(__name__, __file__)
  
apis: numpy.random.uniform numpy.tanh