From bd6d32cbaed087d0fd8130316ad4518af634d45c Mon Sep 17 00:00:00 2001
From: Masahiro Sakai <masahiro.sakai@gmail.com>
Date: Tue, 3 Dec 2019 00:34:22 +0900
Subject: [PATCH] Add test cases for ONNX converter of NStepGRU link

---
 .../functions_tests/test_rnn.py               | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/onnx_chainer_tests/functions_tests/test_rnn.py b/tests/onnx_chainer_tests/functions_tests/test_rnn.py
index 645bf37569e9..8fc2dd094cdd 100644
--- a/tests/onnx_chainer_tests/functions_tests/test_rnn.py
+++ b/tests/onnx_chainer_tests/functions_tests/test_rnn.py
@@ -1,5 +1,6 @@
 import chainer
 import chainer.functions as F
+import chainer.links as L
 from chainer import testing
 
 from onnx_chainer.testing import input_generator
@@ -44,3 +45,33 @@ def __call__(self, hx, ws1, ws2, ws3, bs, xs):
         xs = input_generator.increasing(seq_length, batch_size, input_size)
 
         self.expect(model, (hx, ws1, ws2, ws3, bs, xs))
+
+
+@testing.parameterize(
+    {'n_layers': 1, 'name': 'TestNStepGRU_1_layer'},
+    {'n_layers': 2, 'name': 'TestNStepGRU_2_layer'},
+)
+class TestNStepGRULink(ONNXModelTest):
+    def test_output(self):
+        n_layers = self.n_layers
+        dropout_ratio = 0.0
+        batch_size = 3
+        input_size = 4
+        hidden_size = 5
+        seq_length = 6
+
+        class Model(chainer.Chain):
+            def __init__(self):
+                super().__init__()
+                with self.init_scope():
+                    self.gru = L.NStepGRU(
+                        n_layers, input_size, hidden_size, dropout_ratio)
+
+            def __call__(self, *xs):
+                hy, ys = self.gru(None, xs)
+                return [hy] + ys
+
+        model = Model()
+        xs = [input_generator.increasing(seq_length, input_size)
+              for i in range(batch_size)]
+        self.expect(model, xs, skip_opset_version=[7, 8])