add gru runtime test

Author: masahi

tests/python/relay/test_backend_graph_runtime.py

@@ -7,6 +7,7 @@
 from tvm.relay.scope_builder import ScopeBuilder
 from tvm.relay.op import add
 from tvm.relay.module import Module
+from tvm.relay.testing.config import ctx_list
 
 # @tq, @jr should we put this in testing ns?
 def check_rts(expr, args, expected_result, mod=None):
@@ -127,9 +128,47 @@ def test_plan_memory():
     assert len(device_types) == 1
 
 
+def test_gru():
+    def gru(rnn_dim):
+        X = relay.var("X", shape=(1, rnn_dim))
+        W = relay.var("y", shape=(3 * rnn_dim, rnn_dim))
+        matmul = relay.nn.dense(X, W)
+        splitted = relay.split(matmul, indices_or_sections=3, axis=1)
+        out = relay.sigmoid(splitted[0]) + relay.tanh(splitted[1]) * relay.exp(splitted[2])
+        return relay.Function([X, W], out)
+
+    def sigmoid(x):
+        return 1 / (1 + np.exp(-x))
+
+    def gru_numpy(X, W):
+        prod = np.dot(X, W.transpose())
+        splits = np.split(prod, indices_or_sections=3, axis=1)
+        return sigmoid(splits[0]) + np.tanh(splits[1]) * np.exp(splits[2])
+
+    dtype = "float32"
+    rnn_dim = 1000
+    x = np.random.rand(1, rnn_dim).astype(dtype)
+    y = np.random.rand(3*rnn_dim, rnn_dim).astype(dtype) * 0.01 - 0.005
+    out_shape = (1, rnn_dim)
+    z = gru(rnn_dim)
+
+    for target, ctx in ctx_list():
+        with relay.build_config(opt_level=2):
+            graph, lib, params = relay.build(z, target)
+            m = graph_runtime.create(graph, lib, ctx)
+            m.set_input("X", tvm.nd.array(x.astype(dtype)))
+            m.set_input("y", tvm.nd.array(y.astype(dtype)))
+            m.set_input(**params)
+            m.run()
+            out = m.get_output(0, tvm.nd.empty(out_shape, dtype)).asnumpy()
+            ref = gru_numpy(x, y)
+            tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
 if __name__ == "__main__":
     test_plan_memory()
     test_with_params()
     test_add_op_scalar()
     test_add_op_tensor()
     test_add_op_broadcast()
+    test_gru()