From b6893017d964f531e7f1600faaec4e3817fe89e2 Mon Sep 17 00:00:00 2001
From: notoraptor <stevenbocco@gmail.com>
Date: Thu, 27 Feb 2020 15:53:52 -0500
Subject: [PATCH] Update tests.

---
 tests/python/unittest/test_testing.py |  1 +
 topi/tests/python/test_topi_math.py   | 46 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/tests/python/unittest/test_testing.py b/tests/python/unittest/test_testing.py
index ecf520d251f15..cfa13845e4e8e 100644
--- a/tests/python/unittest/test_testing.py
+++ b/tests/python/unittest/test_testing.py
@@ -31,6 +31,7 @@ def test_check_numerical_grads():
         lambda x: (np.sign(np.sin(1/x)), np.zeros_like(x)),
         lambda x: (x*np.sin(1/x), np.sin(1/x) - np.cos(1/x)/x),
         lambda x: (np.sin(1/x), - np.cos(1/x)/(x*x)),
+        lambda x: (np.tan(x), 1.0 / (np.cos(x) * np.cos(x))),
     ]
 
     # Avoid values too close to 0 since singularities of our functions are there
diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py
index 3a48558ddd529..51ce3a93fa18e 100644
--- a/topi/tests/python/test_topi_math.py
+++ b/topi/tests/python/test_topi_math.py
@@ -133,6 +133,52 @@ def check_device(device):
     test_isnan(-100, 100)
 
 
+def test_ewise_tan():
+    def test_apply(
+        func,
+        name,
+        f_numpy,
+        low,
+        high,
+        shape=(20, 3),
+        dtype='float32',
+        check_round=False,
+        skip_name_check=False,
+    ):
+        A = te.placeholder(dtype=dtype, name="A", shape=shape)
+
+        B = func(A)
+        assert tuple(B.shape) == tuple(A.shape)
+        if not skip_name_check:
+            assert B.op.body[0].name == name
+        a_np = np.random.uniform(low=low, high=high, size=shape).astype(A.dtype) * 10
+        # avoid round check too close to boundary
+        if check_round:
+            a_np += ((np.abs(np.fmod(a_np, 1)) - 0.5) < 1e-6) * 1e-5
+        b_np = f_numpy(a_np)
+
+        def check_device(device):
+            ctx = tvm.context(device, 0)
+            if not ctx.exist:
+                print("Skip because %s is not enabled" % device)
+                return
+            print("Running on target: %s" % device)
+            with tvm.target.create(device):
+                s = topi.testing.get_injective_schedule(device)(B)
+            foo = tvm.build(s, [A, B], device, name=name)
+            a = tvm.nd.array(a_np, ctx)
+            b = tvm.nd.array(np.zeros_like(b_np), ctx)
+            foo(a, b)
+            tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-5, atol=1e-5)
+
+        for target in get_all_backend():
+            check_device(target)
+
+    test_apply(topi.tan, "tan", np.tan, -2.0*np.pi, 2.0*np.pi, dtype='float64')
+    test_apply(topi.tan, "tan", np.tan, -2.0*np.pi, 2.0*np.pi, dtype='float32')
+    test_apply(topi.tan, "tan", np.tan, -2.0*np.pi, 2.0*np.pi, dtype='float16')
+
+
 def test_cast():
     def verify(from_dtype, to_dtype, low=-100, high=100):
         shape = (5, 4)