From cfadf2b7d6be431ec9bc2885d7deb41c6bcc5f52 Mon Sep 17 00:00:00 2001
From: Ramana Radhakrishnan <ramana.radhakrishnan@arm.com>
Date: Sun, 4 Jul 2021 00:23:50 +0100
Subject: [PATCH 1/5] Prepare the topi tests for AArch64 CI.

This pull request cleans up the TOPI testuite for use on the AArch64
CI target by doing the following:

- Introducing a script to run the tests on AArch64 with a suitable
  invocation of the llvm target string by setting the TVM_TEST_TARGETS
  environment variable.

- Cleaning up the use of hard coded targets and moving the testsuite
  to testing more sensibly with the use of tvm.testing.enabled_targets.

- Cleanup the use of tvm.target.create.

- The above allows for the use of tests reasonably with the topi
  tests and cleans up what is needed from the testsuite.

Putting this up for a test run on ci_gpu and ci_cpu to see the effects
of moving TOPI test runs to AArch64 CPU before firing up the Jenkins
changes.

The motivation was from #8361 to pipeclean and add this support.
---
 .../topi/python/test_topi_batch_matmul.py     | 10 ++--
 .../python/test_topi_batch_to_space_nd.py     |  2 +-
 .../topi/python/test_topi_conv2d_NCHWc.py     | 21 +++-----
 .../topi/python/test_topi_conv2d_int8.py      |  6 +--
 .../python/test_topi_conv2d_nhwc_pack_int8.py | 19 ++++---
 .../python/test_topi_deformable_conv2d.py     | 19 ++++---
 tests/python/topi/python/test_topi_lrn.py     | 17 +++---
 tests/python/topi/python/test_topi_reorg.py   | 17 +++---
 .../python/test_topi_space_to_batch_nd.py     |  2 +-
 tests/python/topi/python/test_topi_sparse.py  | 53 ++++++-------------
 .../python/topi/python/test_topi_transform.py |  2 +-
 tests/scripts/task_python_topi_aarch64.sh     | 23 ++++++++
 12 files changed, 88 insertions(+), 103 deletions(-)
 create mode 100755 tests/scripts/task_python_topi_aarch64.sh

diff --git a/tests/python/topi/python/test_topi_batch_matmul.py b/tests/python/topi/python/test_topi_batch_matmul.py
index 8c8ad37287dc..fcddc748aaf2 100644
--- a/tests/python/topi/python/test_topi_batch_matmul.py
+++ b/tests/python/topi/python/test_topi_batch_matmul.py
@@ -110,14 +110,12 @@ def get_ref_data():
     # get the test data
     a_np, b_np, c_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
+    def check_device(target, dev):
         if device == "cuda" and not tvm.contrib.nvcc.have_int8(dev.compute_version):
             print("Skip because int8 intrinsics are not available")
             return
 
-        print("Running on target: %s" % device)
-        with tvm.target.Target(device):
+        with tvm.target.Target(target):
             out = topi.cuda.batch_matmul_int8(x, y, None, out_dtype)
             s = topi.cuda.schedule_batch_matmul_int8([out])
         a = tvm.nd.array(a_np, dev)
@@ -127,8 +125,8 @@ def check_device(device):
         f(a, b, c)
         tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
 
-    for device in ["cuda"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 @tvm.testing.uses_gpu
diff --git a/tests/python/topi/python/test_topi_batch_to_space_nd.py b/tests/python/topi/python/test_topi_batch_to_space_nd.py
index 6ee99dacc61a..d87d4b9987cc 100644
--- a/tests/python/topi/python/test_topi_batch_to_space_nd.py
+++ b/tests/python/topi/python/test_topi_batch_to_space_nd.py
@@ -44,7 +44,7 @@ def verify_batch_to_space_nd(input_shape, block_shape, crop_begin_list, crop_end
 
     def check_device(target, dev):
         print("Running on target: %s" % target)
-        with tvm.target.create(target):
+        with tvm.target.Target(target):
             s = tvm.topi.testing.get_injective_schedule(target)(B)
         a = tvm.nd.array(a_np, dev)
         b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), dev)
diff --git a/tests/python/topi/python/test_topi_conv2d_NCHWc.py b/tests/python/topi/python/test_topi_conv2d_NCHWc.py
index 2298816d373a..561e9ced07cb 100644
--- a/tests/python/topi/python/test_topi_conv2d_NCHWc.py
+++ b/tests/python/topi/python/test_topi_conv2d_NCHWc.py
@@ -51,7 +51,7 @@ def _transform_bias(bias, bn):
     bias = np.transpose(bias, (0, 2, 3, 1))
     return bias
 
-
+@tvm.testing.requires_llvm
 def verify_conv2d_NCHWc(
     batch,
     in_channel,
@@ -115,13 +115,8 @@ def get_ref_data():
 
     a_np, w_np, b_np, c_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
-        with tvm.target.Target(device):
+    def check_device(target, dev):
+        with tvm.target.Target(target):
             C = topi.x86.conv2d_NCHWc(
                 A,
                 W,
@@ -146,7 +141,7 @@ def check_device(device):
             func = tvm.build(
                 s,
                 [A, W, bias, C],
-                device,
+                target,
                 name="relu_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation),
             )
@@ -155,7 +150,7 @@ def check_device(device):
             func = tvm.build(
                 s,
                 [A, W, C],
-                device,
+                target,
                 name="relu_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation),
             )
@@ -163,9 +158,9 @@ def check_device(device):
         tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
 
     # test llvm only for now since conv2d_NCHWc implement is missing in other backend.
-    for device in ["llvm"]:
-        with autotvm.tophub.context(device):  # load tophub pre-tuned parameters
-            check_device(device)
+    for target, device in tvm.testing.enabled_targets():
+        with autotvm.tophub.context(target):  # load tophub pre-tuned parameters
+            check_device(target, device)
 
 
 def test_conv2d_NCHWc():
diff --git a/tests/python/topi/python/test_topi_conv2d_int8.py b/tests/python/topi/python/test_topi_conv2d_int8.py
index 66d280e85c75..08f94184f84a 100644
--- a/tests/python/topi/python/test_topi_conv2d_int8.py
+++ b/tests/python/topi/python/test_topi_conv2d_int8.py
@@ -175,8 +175,7 @@ def get_ref_data():
 
     a_np, w_np, b_np, c_np = get_ref_data()
 
-    def check_target(target):
-        dev = tvm.device(target, 0)
+    def check_target(target, dev):
         if not tvm.testing.device_enabled(target):
             print("Skip because %s is not enabled" % target)
             return
@@ -222,7 +221,8 @@ def check_target(target):
             func(a, w, c)
         tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
 
-    check_target("llvm")
+    for target, dev in tvm.testing.enabled_targets():
+        check_target(target, dev)
 
 
 oc_block_factor = 4
diff --git a/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py b/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
index 8b20961a8cdf..8ea6a238d9e6 100644
--- a/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
+++ b/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
@@ -26,7 +26,7 @@
 import tvm.topi.testing
 from tvm.contrib.pickle_memoize import memoize
 from tvm.topi.utils import get_const_tuple
-
+import tvm.testing
 
 def verify_conv2d_1x1_nhwc_pack_int8(
     batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1
@@ -51,26 +51,25 @@ def get_ref_data():
 
     a_np, w_np, b_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
+    def check_device(target, dev):
+        if not tvm.testing.device_enabled(target):
+            print("Skip because %s is not enabled" % target)
             return
-        print("Running on target: %s" % device)
+        print("Running on target: %s" % target)
 
-        with tvm.target.Target(device):
+        with tvm.target.Target(target):
             B = topi.nn.conv2d(A, W, stride, padding, dilation, layout="NHWC", out_dtype="int32")
             s = topi.x86.schedule_conv2d_nhwc_pack_int8([B])
         a = tvm.nd.array(a_np, dev)
         w = tvm.nd.array(w_np, dev)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
-        func = tvm.build(s, [A, W, B], device)
+        func = tvm.build(s, [A, W, B], target)
         func(a, w, b)
         tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
 
     # for device in ['llvm -mcpu=skylake-avx512']:
-    for device in ["llvm"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 # TODO(@llyfacebook): Please fix https://github.com/apache/tvm/issues/4122 to enable this test.
diff --git a/tests/python/topi/python/test_topi_deformable_conv2d.py b/tests/python/topi/python/test_topi_deformable_conv2d.py
index 70cc9a690cdc..fdfe9974804c 100644
--- a/tests/python/topi/python/test_topi_deformable_conv2d.py
+++ b/tests/python/topi/python/test_topi_deformable_conv2d.py
@@ -92,14 +92,13 @@ def get_ref_data():
 
     a_np, offset_np, w_np, c_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
+    def check_device(target, dev):
+        if not tvm.testing.device_enabled(target):
+            print("Skip because %s is not enabled" % target)
             return
-        print("Running on target: %s" % device)
-        fcompute, fschedule = tvm.topi.testing.dispatch(device, _deformable_conv2d_nchw_implement)
-        with tvm.target.Target(device):
+        print("Running on target: %s" % target)
+        fcompute, fschedule = tvm.topi.testing.dispatch(target, _deformable_conv2d_nchw_implement)
+        with tvm.target.Target(target):
             C = fcompute(A, Offset, W, stride, padding, dilation, deformable_groups, groups, dtype)
             s = fschedule([C])
 
@@ -108,12 +107,12 @@ def check_device(device):
             w = tvm.nd.array(w_np, dev)
             c = tvm.nd.empty(c_np.shape, dtype=c_np.dtype, device=dev)
 
-            func = tvm.build(s, [A, Offset, W, C], device)
+            func = tvm.build(s, [A, Offset, W, C], target)
             func(a, offset, w, c)
             tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
 
-    for device in ["llvm", "cuda"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 def verify_deformable_conv2d_nhwc(
diff --git a/tests/python/topi/python/test_topi_lrn.py b/tests/python/topi/python/test_topi_lrn.py
index f9fb7dbd4ec4..6619f50d0155 100644
--- a/tests/python/topi/python/test_topi_lrn.py
+++ b/tests/python/topi/python/test_topi_lrn.py
@@ -42,23 +42,18 @@ def verify_lrn(shape, size, axis, bias, alpha, beta):
     a_np = np.random.uniform(size=shape).astype(dtype)
     b_np = tvm.topi.testing.lrn_python(a_np, size, axis, bias, alpha, beta)
 
-    def check_device(device):
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
-        with tvm.target.Target(device):
-            s_func = tvm.topi.testing.dispatch(device, _lrn_schedule)
+    def check_device(target, dev):
+        with tvm.target.Target(target):
+            s_func = tvm.topi.testing.dispatch(target, _lrn_schedule)
             s = s_func([B])
-        dev = tvm.device(device, 0)
         a = tvm.nd.array(a_np, dev)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), dev)
-        f = tvm.build(s, [A, B], device)
+        f = tvm.build(s, [A, B], target)
         f(a, b)
         tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
 
-    for device in ["llvm", "cuda", "opencl", "metal", "rocm", "vulkan", "nvptx"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 @tvm.testing.uses_gpu
diff --git a/tests/python/topi/python/test_topi_reorg.py b/tests/python/topi/python/test_topi_reorg.py
index f41b4b740bec..d2f01109e2b8 100644
--- a/tests/python/topi/python/test_topi_reorg.py
+++ b/tests/python/topi/python/test_topi_reorg.py
@@ -46,24 +46,19 @@ def get_ref_data_reorg():
 
     a_np, b_np = get_ref_data_reorg()
 
-    def check_device(device):
+    def check_device(target, dev):
         """Cheching devices is enabled or not"""
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
-        with tvm.target.Target(device):
-            s_func = tvm.topi.testing.dispatch(device, _reorg_schedule)
+        with tvm.target.Target(target):
+            s_func = tvm.topi.testing.dispatch(target, _reorg_schedule)
             s = s_func([B])
         a = tvm.nd.array(a_np, dev)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
-        func = tvm.build(s, [A, B], device)
+        func = tvm.build(s, [A, B], target)
         func(a, b)
         tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
 
-    for device in ["llvm", "cuda"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 @tvm.testing.uses_gpu
diff --git a/tests/python/topi/python/test_topi_space_to_batch_nd.py b/tests/python/topi/python/test_topi_space_to_batch_nd.py
index 039f91aa059e..d0aaacce0c71 100644
--- a/tests/python/topi/python/test_topi_space_to_batch_nd.py
+++ b/tests/python/topi/python/test_topi_space_to_batch_nd.py
@@ -44,7 +44,7 @@ def verify_space_to_batch_nd(input_shape, block_shape, pad_before, pad_after, pa
 
     def check_target(target, dev):
         print("Running on target: %s" % target)
-        with tvm.target.create(target):
+        with tvm.target.Target(target):
             s = tvm.topi.testing.get_injective_schedule(target)(B)
         a = tvm.nd.array(a_np, dev)
         b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), dev)
diff --git a/tests/python/topi/python/test_topi_sparse.py b/tests/python/topi/python/test_topi_sparse.py
index 11006576fea3..e273050725f4 100644
--- a/tests/python/topi/python/test_topi_sparse.py
+++ b/tests/python/topi/python/test_topi_sparse.py
@@ -57,12 +57,7 @@ def get_ref_data():
 
     a_np, b_np, c_np, d_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
+    def check_device(target, dev):
         a = tvmsp.array(a_np, dev)
         _nr, _nc, _n = a.shape[0], a.shape[1], a.data.shape[0]
         assert a.shape[0] == a.indptr.shape[0] - 1
@@ -72,12 +67,12 @@ def check_device(device):
         assert a.data.dtype == A.data.dtype
         assert a.indices.dtype == A.indices.dtype
         assert a.indptr.dtype == A.indptr.dtype
-        f = tvm.build(s, [nr, A.data, A.indices, A.indptr, B, C, D], device, name="csrmv")
+        f = tvm.build(s, [nr, A.data, A.indices, A.indptr, B, C, D], target, name="csrmv")
         f(_nr, a.data, a.indices, a.indptr, b, c, d)
         tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-4, atol=1e-4)
 
-    for device in ["llvm"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 def verify_dynamic_csrmm(batch, in_dim, out_dim, dtype, use_bias=True):
@@ -102,25 +97,20 @@ def get_ref_data():
 
     a_np, b_np, c_np, d_np = get_ref_data()
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
+    def check_device(target, dev):
         a = tvmsp.array(a_np, dev)
         _nr, _nc, _n = a.shape[0], a.shape[1], a.data.shape[0]
         assert a.shape[0] == a.indptr.shape[0] - 1
         b = tvm.nd.array(b_np, dev)
         c = tvm.nd.array(c_np, dev)
         d = tvm.nd.array(np.zeros((_nr, out_dim), dtype=dtype), dev)
-        f = tvm.build(s, [nr, A.data, A.indices, A.indptr, B, C, D], device, name="csrmm")
+        f = tvm.build(s, [nr, A.data, A.indices, A.indptr, B, C, D], target, name="csrmm")
 
         f(_nr, a.data, a.indices, a.indptr, b, c, d)
         tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-2, atol=1e-2)
 
-    for device in ["llvm"]:
-        check_device(device)
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 def verify_dense_si(batch, in_dim, out_dim, use_bias=True, dtype="float32"):
@@ -437,14 +427,9 @@ def test_sparse_dense_bsr_randomized():
         W_indptr = te.placeholder(shape=W_sp_np.indptr.shape, dtype=str(W_sp_np.indptr.dtype))
         X = te.placeholder(shape=X_np.shape, dtype=str(X_np.dtype))
 
-        def check_device(device):
-            dev = tvm.device(device, 0)
-            if not tvm.testing.device_enabled(device):
-                print("Skip because %s is not enabled" % device)
-                return
-            print("Running on target: %s" % device)
-            fcompute, fschedule = tvm.topi.testing.dispatch(device, _sparse_dense_implement)
-            with tvm.target.Target(device):
+        def check_device(target, dev):
+            fcompute, fschedule = tvm.topi.testing.dispatch(target, _sparse_dense_implement)
+            with tvm.target.Target(target):
                 Y = fcompute(X, W_data, W_indices, W_indptr)
                 s = fschedule([Y])
                 func = tvm.build(s, [X, W_data, W_indices, W_indptr, Y])
@@ -458,8 +443,8 @@ def check_device(device):
                 )
                 tvm.testing.assert_allclose(Y_tvm.numpy(), Y_np, atol=1e-5, rtol=1e-5)
 
-        for device in ["llvm", "cuda"]:
-            check_device(device)
+        for target, dev in tvm.testing.enabled_targets():
+            check_device(target, dev)
 
 
 @tvm.testing.parametrize_targets("cuda", "rocm")
@@ -576,14 +561,9 @@ def verify_sparse_conv2d_bsr(M, H, W, N, K, BS_R, BS_C, density, layout):
     Y = topi.nn.sparse_conv2d(X, W_data, W_indices, W_indptr, layout)
     s = te.create_schedule(Y.op)
 
-    def check_device(device):
-        dev = tvm.device(device, 0)
-        if not tvm.testing.device_enabled(device):
-            print("Skip because %s is not enabled" % device)
-            return
-        print("Running on target: %s" % device)
+    def check_device(target, dev):
 
-        func = tvm.build(s, [X, W_data, W_indices, W_indptr, Y])
+        func = tvm.build(s, [X, W_data, W_indices, W_indptr, Y], target)
         Y_tvm = tvm.nd.array(np.zeros(Y_np.shape, dtype="float32"))
         func(
             tvm.nd.array(X_np, dev),
@@ -594,7 +574,8 @@ def check_device(device):
         )
         tvm.testing.assert_allclose(Y_tvm.numpy(), Y_np.astype("float32"), atol=1e-4, rtol=1e-4)
 
-    check_device("llvm")
+    for target, dev in tvm.testing.enabled_targets():
+        check_device(target, dev)
 
 
 def test_sparse_conv2d_bsr():
diff --git a/tests/python/topi/python/test_topi_transform.py b/tests/python/topi/python/test_topi_transform.py
index ddde2e20e754..701513c33139 100644
--- a/tests/python/topi/python/test_topi_transform.py
+++ b/tests/python/topi/python/test_topi_transform.py
@@ -794,7 +794,7 @@ def check_device(target, dev):
             print("Skip because %s is not enabled" % target)
             return
         print("Running on target: %s" % target)
-        with tvm.target.create(target):
+        with tvm.target.Target(target):
             s = tvm.topi.testing.get_injective_schedule(target)(out)
 
         func = tvm.build(s, [data] + indices + [out], target, name="adv_index")
diff --git a/tests/scripts/task_python_topi_aarch64.sh b/tests/scripts/task_python_topi_aarch64.sh
new file mode 100755
index 000000000000..c809dbf15f05
--- /dev/null
+++ b/tests/scripts/task_python_topi_aarch64.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+set -e
+set -u
+
+export TVM_TEST_TARGETS="llvm -device=arm_cpu"
+./tests/scripts/task_python_topi.sh

From 0828a584d60b13a901ddb0173742ac8e2bcbec96 Mon Sep 17 00:00:00 2001
From: Ramana Radhakrishnan <ramana.radhakrishnan@arm.com>
Date: Mon, 5 Jul 2021 23:44:28 +0100
Subject: [PATCH 2/5] Fix use of target.create vs target.Target

Change-Id: I70537a20bc683c490ef3a504a14136e7b0a1456d
---
 .../topi/python/test_topi_conv2d_nhwc.py       | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tests/python/topi/python/test_topi_conv2d_nhwc.py b/tests/python/topi/python/test_topi_conv2d_nhwc.py
index eb4c5a343b58..9d4c0b22ba17 100644
--- a/tests/python/topi/python/test_topi_conv2d_nhwc.py
+++ b/tests/python/topi/python/test_topi_conv2d_nhwc.py
@@ -66,21 +66,25 @@ def get_ref_data():
 
     a_np, w_np, b_np = get_ref_data()
 
-    def check_device(target, dev):
-        print("Running on target: %s" % target)
-        with tvm.target.Target(target):
-            fcompute, fschedule = tvm.topi.testing.dispatch(target, _conv2d_nhwc_implement)
+    def check_device(device):
+        if not tvm.testing.device_enabled(device):
+            print("Skip because %s is not enabled" % device)
+            return
+        print("Running on target: %s" % device)
+        with tvm.target.Target(device):
+            fcompute, fschedule = tvm.topi.testing.dispatch(device, _conv2d_nhwc_implement)
             B = fcompute(A, W, stride, padding, dilation, dtype)
             s = fschedule([B])
+        dev = tvm.device(device, 0)
         a = tvm.nd.array(a_np, dev)
         w = tvm.nd.array(w_np, dev)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
-        func = tvm.build(s, [A, W, B], target)
+        func = tvm.build(s, [A, W, B], device)
         func(a, w, b)
         tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
 
-    for target, dev in tvm.testing.enabled_targets():
-        check_device(target, dev)
+    for device in ["llvm", "cuda"]:
+        check_device(device)
 
 
 @tvm.testing.uses_gpu

From c2a6a8608ccb5273a2af81280ff0e2ff2b66d232 Mon Sep 17 00:00:00 2001
From: Ramana Radhakrishnan <ramana.radhakrishnan@arm.com>
Date: Tue, 3 Aug 2021 21:59:43 +0100
Subject: [PATCH 3/5] Fixup error with test_topi_batch_matmul.py

---
 tests/python/topi/python/test_topi_batch_matmul.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/topi/python/test_topi_batch_matmul.py b/tests/python/topi/python/test_topi_batch_matmul.py
index fcddc748aaf2..28039cb5cbb8 100644
--- a/tests/python/topi/python/test_topi_batch_matmul.py
+++ b/tests/python/topi/python/test_topi_batch_matmul.py
@@ -111,7 +111,7 @@ def get_ref_data():
     a_np, b_np, c_np = get_ref_data()
 
     def check_device(target, dev):
-        if device == "cuda" and not tvm.contrib.nvcc.have_int8(dev.compute_version):
+        if target == "cuda" and not tvm.contrib.nvcc.have_int8(dev.compute_version):
             print("Skip because int8 intrinsics are not available")
             return
 

From b91ceb6a3e6c795936ec245e12a7d85a9e5f442b Mon Sep 17 00:00:00 2001
From: Ramana Radhakrishnan <ramana.radhakrishnan@arm.com>
Date: Tue, 3 Aug 2021 22:09:06 +0100
Subject: [PATCH 4/5] Fix topi testing for test_topi_sparse.py

---
 tests/python/topi/python/test_topi_sparse.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tests/python/topi/python/test_topi_sparse.py b/tests/python/topi/python/test_topi_sparse.py
index e273050725f4..35492c9a5f5c 100644
--- a/tests/python/topi/python/test_topi_sparse.py
+++ b/tests/python/topi/python/test_topi_sparse.py
@@ -71,7 +71,10 @@ def check_device(target, dev):
         f(_nr, a.data, a.indices, a.indptr, b, c, d)
         tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-4, atol=1e-4)
 
-    for target, dev in tvm.testing.enabled_targets():
+    # Unable to use the following idiom and thus restrict to llvm only.
+    #for target, dev in tvm.testing.enabled_targets():
+    for target in ['llvm']:
+        dev = tvm.device(target, 0)
         check_device(target, dev)
 
 
@@ -109,7 +112,10 @@ def check_device(target, dev):
         f(_nr, a.data, a.indices, a.indptr, b, c, d)
         tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-2, atol=1e-2)
 
-    for target, dev in tvm.testing.enabled_targets():
+    # This test is not yet ready to use the idiom ,
+    # for target, dev in tvm.testing.enabled_targets()
+    for target in ['llvm']:
+        dev = tvm.device(target, 0)
         check_device(target, dev)
 
 
@@ -574,7 +580,10 @@ def check_device(target, dev):
         )
         tvm.testing.assert_allclose(Y_tvm.numpy(), Y_np.astype("float32"), atol=1e-4, rtol=1e-4)
 
-    for target, dev in tvm.testing.enabled_targets():
+    #Unable to use the idiom below as the test isn't suitable for non llvm targets.
+    #for target, dev in tvm.testing.enabled_targets():
+    for target in ['llvm']:
+        dev = tvm.device(target, 0)
         check_device(target, dev)
 
 

From ffa060c4442da00da30c1b935516bfe4e2d6ac2c Mon Sep 17 00:00:00 2001
From: Ramana Radhakrishnan <ramana.radhakrishnan@arm.com>
Date: Fri, 6 Aug 2021 00:18:30 +0100
Subject: [PATCH 5/5] Fix black formatting

---
 .../python/topi/python/test_topi_argwhere.py  |  5 +---
 .../test_topi_batch_matmul_tensorcore.py      |  2 +-
 .../python/topi/python/test_topi_broadcast.py | 25 +++----------------
 .../topi/python/test_topi_conv2d_NCHWc.py     |  1 +
 .../topi/python/test_topi_conv2d_nchw.py      | 22 +++-------------
 .../topi/python/test_topi_conv2d_nhwc.py      | 10 ++------
 .../python/test_topi_conv2d_nhwc_pack_int8.py |  1 +
 .../topi/python/test_topi_conv2d_winograd.py  | 11 +-------
 .../topi/python/test_topi_conv3d_winograd.py  |  2 +-
 .../python/test_topi_deformable_conv2d.py     |  2 +-
 tests/python/topi/python/test_topi_dense.py   | 19 +++-----------
 .../topi/python/test_topi_depthwise_conv2d.py | 18 ++++---------
 .../topi/python/test_topi_group_conv2d.py     |  2 +-
 tests/python/topi/python/test_topi_math.py    |  7 +-----
 tests/python/topi/python/test_topi_matmul.py  |  4 +--
 tests/python/topi/python/test_topi_relu.py    |  8 ++----
 tests/python/topi/python/test_topi_reorg.py   |  5 +---
 tests/python/topi/python/test_topi_scan.py    |  7 +-----
 tests/python/topi/python/test_topi_sparse.py  | 12 ++++-----
 tests/python/topi/python/test_topi_vision.py  |  5 +---
 20 files changed, 37 insertions(+), 131 deletions(-)

diff --git a/tests/python/topi/python/test_topi_argwhere.py b/tests/python/topi/python/test_topi_argwhere.py
index 8592f57b74a4..f0145482a543 100644
--- a/tests/python/topi/python/test_topi_argwhere.py
+++ b/tests/python/topi/python/test_topi_argwhere.py
@@ -22,10 +22,7 @@
 from tvm import topi
 import tvm.topi.testing
 
-_argwhere_schedule = {
-    "generic": topi.generic.schedule_argwhere,
-    "gpu": topi.cuda.schedule_argwhere,
-}
+_argwhere_schedule = {"generic": topi.generic.schedule_argwhere, "gpu": topi.cuda.schedule_argwhere}
 
 _argwhere_compute = {"llvm": topi.argwhere, "cuda": topi.cuda.argwhere}
 
diff --git a/tests/python/topi/python/test_topi_batch_matmul_tensorcore.py b/tests/python/topi/python/test_topi_batch_matmul_tensorcore.py
index eb657a329889..6190bf5e2c63 100644
--- a/tests/python/topi/python/test_topi_batch_matmul_tensorcore.py
+++ b/tests/python/topi/python/test_topi_batch_matmul_tensorcore.py
@@ -26,7 +26,7 @@
 import tvm.testing
 
 _batch_matmul_implement = {
-    "gpu": (topi.cuda.batch_matmul_tensorcore, topi.cuda.schedule_batch_matmul_tensorcore),
+    "gpu": (topi.cuda.batch_matmul_tensorcore, topi.cuda.schedule_batch_matmul_tensorcore)
 }
 
 
diff --git a/tests/python/topi/python/test_topi_broadcast.py b/tests/python/topi/python/test_topi_broadcast.py
index d77ed2eae2e4..1ee0dc65efba 100644
--- a/tests/python/topi/python/test_topi_broadcast.py
+++ b/tests/python/topi/python/test_topi_broadcast.py
@@ -290,13 +290,7 @@ def test_shift():
 
 @tvm.testing.uses_gpu
 def test_logical_single_ele():
-    def test_apply(
-        func,
-        name,
-        f_numpy,
-        indata,
-        dtype="bool",
-    ):
+    def test_apply(func, name, f_numpy, indata, dtype="bool"):
         # Build the logic and compile the function
         A = te.placeholder(shape=indata.shape, name="A", dtype=dtype)
         B = func(A)
@@ -327,13 +321,7 @@ def check_target(target, dev):
 
 @tvm.testing.uses_gpu
 def test_bitwise_not():
-    def test_apply(
-        func,
-        name,
-        f_numpy,
-        shape,
-        dtype="int32",
-    ):
+    def test_apply(func, name, f_numpy, shape, dtype="int32"):
         # Build the logic and compile the function
         A = te.placeholder(shape=shape, name="A", dtype=dtype)
         B = func(A)
@@ -365,14 +353,7 @@ def check_target(target, dev):
 
 @tvm.testing.uses_gpu
 def test_logical_binary_ele():
-    def test_apply(
-        func,
-        name,
-        f_numpy,
-        lhs,
-        rhs,
-        dtype="bool",
-    ):
+    def test_apply(func, name, f_numpy, lhs, rhs, dtype="bool"):
         # Build the logic and compile the function
         A = te.var("A", dtype=dtype)
         B = te.var("B", dtype=dtype)
diff --git a/tests/python/topi/python/test_topi_conv2d_NCHWc.py b/tests/python/topi/python/test_topi_conv2d_NCHWc.py
index 561e9ced07cb..ea2529f7fe0c 100644
--- a/tests/python/topi/python/test_topi_conv2d_NCHWc.py
+++ b/tests/python/topi/python/test_topi_conv2d_NCHWc.py
@@ -51,6 +51,7 @@ def _transform_bias(bias, bn):
     bias = np.transpose(bias, (0, 2, 3, 1))
     return bias
 
+
 @tvm.testing.requires_llvm
 def verify_conv2d_NCHWc(
     batch,
diff --git a/tests/python/topi/python/test_topi_conv2d_nchw.py b/tests/python/topi/python/test_topi_conv2d_nchw.py
index 2a4865c6dd8d..2f3cb1bffd96 100644
--- a/tests/python/topi/python/test_topi_conv2d_nchw.py
+++ b/tests/python/topi/python/test_topi_conv2d_nchw.py
@@ -52,15 +52,7 @@ def bias_shape(num_filter):
 
 @tvm.testing.fixture(cache_return_value=True)
 def ref_data(
-    input_shape,
-    weight_shape,
-    bias_shape,
-    dtype,
-    stride,
-    padding,
-    dilation,
-    add_bias,
-    apply_relu,
+    input_shape, weight_shape, bias_shape, dtype, stride, padding, dilation, add_bias, apply_relu
 ):
     a_np = np.random.uniform(size=input_shape).astype(dtype)
     w_np = np.random.uniform(size=weight_shape).astype(dtype)
@@ -146,15 +138,7 @@ def test_conv2d_nchw(
 
     @tvm.testing.parametrize_targets("llvm")
     def test_workload_padding(
-        self,
-        target,
-        input_shape,
-        weight_shape,
-        stride,
-        padding,
-        dilation,
-        dtype,
-        ref_data,
+        self, target, input_shape, weight_shape, stride, padding, dilation, dtype, ref_data
     ):
         a_np, w_np, b_np, c_np = ref_data
         _, _, out_height, out_width = c_np.shape
@@ -282,7 +266,7 @@ class TestAsymmetricPadding(BaseConv2DTests):
 
 class TestBatchSize(BaseConv2DTests):
     in_channel, in_size, num_filter, kernel, stride, padding = tvm.testing.parameters(
-        (64, 56, 64, 3, 1, 1),
+        (64, 56, 64, 3, 1, 1)
     )
     batch = tvm.testing.parameter(1, 4, 9)
 
diff --git a/tests/python/topi/python/test_topi_conv2d_nhwc.py b/tests/python/topi/python/test_topi_conv2d_nhwc.py
index 9d4c0b22ba17..799e86155029 100644
--- a/tests/python/topi/python/test_topi_conv2d_nhwc.py
+++ b/tests/python/topi/python/test_topi_conv2d_nhwc.py
@@ -34,14 +34,8 @@
         topi.arm_cpu.conv2d_nhwc_spatial_pack,
         topi.arm_cpu.schedule_conv2d_nhwc_spatial_pack,
     ),
-    "mali": (
-        topi.mali.conv2d_nhwc_spatial_pack,
-        topi.mali.schedule_conv2d_nhwc_spatial_pack,
-    ),
-    "bifrost": (
-        topi.mali.conv2d_nhwc_spatial_pack,
-        topi.mali.schedule_conv2d_nhwc_spatial_pack,
-    ),
+    "mali": (topi.mali.conv2d_nhwc_spatial_pack, topi.mali.schedule_conv2d_nhwc_spatial_pack),
+    "bifrost": (topi.mali.conv2d_nhwc_spatial_pack, topi.mali.schedule_conv2d_nhwc_spatial_pack),
     "hls": (topi.nn.conv2d_nhwc, topi.hls.schedule_conv2d_nhwc),
 }
 
diff --git a/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py b/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
index 8ea6a238d9e6..2000ed5a49c3 100644
--- a/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
+++ b/tests/python/topi/python/test_topi_conv2d_nhwc_pack_int8.py
@@ -28,6 +28,7 @@
 from tvm.topi.utils import get_const_tuple
 import tvm.testing
 
+
 def verify_conv2d_1x1_nhwc_pack_int8(
     batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1
 ):
diff --git a/tests/python/topi/python/test_topi_conv2d_winograd.py b/tests/python/topi/python/test_topi_conv2d_winograd.py
index 82368f118f32..0a739deda06e 100644
--- a/tests/python/topi/python/test_topi_conv2d_winograd.py
+++ b/tests/python/topi/python/test_topi_conv2d_winograd.py
@@ -170,16 +170,7 @@ def test_conv2d_nchw():
     verify_conv2d_nchw(1, 48, 35, 48, 5, 1, "VALID", devices=["cuda"])
 
 
-def verify_conv2d_nhwc(
-    batch,
-    in_channel,
-    in_size,
-    num_filter,
-    kernel,
-    stride,
-    padding,
-    dilation=1,
-):
+def verify_conv2d_nhwc(batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1):
     # This version is intented to be used by the auto-scheduler,
     # so we only test the correctness of compute declaration
     # with the default naive schedule in cpu
diff --git a/tests/python/topi/python/test_topi_conv3d_winograd.py b/tests/python/topi/python/test_topi_conv3d_winograd.py
index 0b9d579287c3..6456600b9049 100644
--- a/tests/python/topi/python/test_topi_conv3d_winograd.py
+++ b/tests/python/topi/python/test_topi_conv3d_winograd.py
@@ -29,7 +29,7 @@
 
 
 _conv3d_ncdhw_implement = {
-    "gpu": (topi.cuda.conv3d_ncdhw_winograd, topi.cuda.schedule_conv3d_ncdhw_winograd),
+    "gpu": (topi.cuda.conv3d_ncdhw_winograd, topi.cuda.schedule_conv3d_ncdhw_winograd)
 }
 
 
diff --git a/tests/python/topi/python/test_topi_deformable_conv2d.py b/tests/python/topi/python/test_topi_deformable_conv2d.py
index fdfe9974804c..a3938715e9c2 100644
--- a/tests/python/topi/python/test_topi_deformable_conv2d.py
+++ b/tests/python/topi/python/test_topi_deformable_conv2d.py
@@ -32,7 +32,7 @@
 }
 
 _deformable_conv2d_nhwc_implement = {
-    "generic": (topi.nn.deformable_conv2d_nhwc, topi.generic.schedule_deformable_conv2d_nhwc),
+    "generic": (topi.nn.deformable_conv2d_nhwc, topi.generic.schedule_deformable_conv2d_nhwc)
 }
 
 
diff --git a/tests/python/topi/python/test_topi_dense.py b/tests/python/topi/python/test_topi_dense.py
index 964c1621fa47..fc8e02edd0e7 100644
--- a/tests/python/topi/python/test_topi_dense.py
+++ b/tests/python/topi/python/test_topi_dense.py
@@ -31,10 +31,7 @@
 use_bias = tvm.testing.parameter(True, False)
 batch_size = tvm.testing.parameter(1, 2, 128)
 in_dim, out_dim = tvm.testing.parameters((1024, 1000))
-in_dtype, out_dtype = tvm.testing.parameters(
-    ("float32", "float32"),
-    ("int8", "int32"),
-)
+in_dtype, out_dtype = tvm.testing.parameters(("float32", "float32"), ("int8", "int32"))
 
 
 _dense_implementations = {
@@ -137,19 +134,9 @@ def test_dense(
 @pytest.mark.parametrize("target,in_dtype,out_dtype", [("cuda", "int8", "int32")])
 @tvm.testing.requires_gpu
 def test_dense_cuda_int8(
-    target,
-    dev,
-    batch_size,
-    in_dim,
-    out_dim,
-    use_bias,
-    dense_ref_data,
-    in_dtype,
-    out_dtype,
+    target, dev, batch_size, in_dim, out_dim, use_bias, dense_ref_data, in_dtype, out_dtype
 ):
-    implementations = [
-        (topi.cuda.dense_int8, topi.cuda.schedule_dense_int8),
-    ]
+    implementations = [(topi.cuda.dense_int8, topi.cuda.schedule_dense_int8)]
     with Int8Fallback():
         test_dense(
             target,
diff --git a/tests/python/topi/python/test_topi_depthwise_conv2d.py b/tests/python/topi/python/test_topi_depthwise_conv2d.py
index 092ac9df5f9a..58ebebfe66fe 100644
--- a/tests/python/topi/python/test_topi_depthwise_conv2d.py
+++ b/tests/python/topi/python/test_topi_depthwise_conv2d.py
@@ -63,7 +63,7 @@
         "gpu": [(topi.nn.depthwise_conv2d_nhwc, topi.cuda.schedule_depthwise_conv2d_nhwc)],
     },
     "NCHWc": {
-        "generic": [(topi.x86.depthwise_conv2d_NCHWc, topi.x86.schedule_depthwise_conv2d_NCHWc)],
+        "generic": [(topi.x86.depthwise_conv2d_NCHWc, topi.x86.schedule_depthwise_conv2d_NCHWc)]
     },
 }
 
@@ -158,13 +158,7 @@ def ref_data(
     if apply_relu:
         output_np = np.maximum(output_np, 0)
 
-    return (
-        input_np,
-        filter_np,
-        scale_np,
-        shift_np,
-        output_np,
-    )
+    return (input_np, filter_np, scale_np, shift_np, output_np)
 
 
 class BaseDepthwiseConv2D:
@@ -178,8 +172,7 @@ class BaseDepthwiseConv2D:
     layout = tvm.testing.parameter("NCHW", "NHWC")
 
     (batch, in_channel, in_size, channel_multiplier, kernel, stride) = tvm.testing.parameters(
-        (1, 728, 32, 1, 3, 1),
-        (4, 256, 64, 2, 5, 2),
+        (1, 728, 32, 1, 3, 1), (4, 256, 64, 2, 5, 2)
     )
     padding = tvm.testing.parameter("SAME", "VALID")
     dilation = tvm.testing.parameter(1, 2)
@@ -280,8 +273,7 @@ def test_conv2d(
                     scale_tvm = tvm.nd.array(scale_np, dev)
                     shift_tvm = tvm.nd.array(shift_np, dev)
                     output_tvm = tvm.nd.array(
-                        np.zeros(shape=get_const_tuple(C.shape), dtype=C.dtype),
-                        dev,
+                        np.zeros(shape=get_const_tuple(C.shape), dtype=C.dtype), dev
                     )
 
                     f(input_tvm, filter_tvm, scale_tvm, shift_tvm, output_tvm)
@@ -361,7 +353,7 @@ class TestDepthwiseConv2D_NCHWc(BaseDepthwiseConv2D):
     # depthwise_conv2d_NCHWc currently does not support channel multiplier > 1
     layout = tvm.testing.parameter("NCHWc")
     (batch, in_channel, in_size, channel_multiplier, kernel, stride) = tvm.testing.parameters(
-        (1, 728, 32, 1, 3, 1),
+        (1, 728, 32, 1, 3, 1)
     )
 
 
diff --git a/tests/python/topi/python/test_topi_group_conv2d.py b/tests/python/topi/python/test_topi_group_conv2d.py
index 55b24feece93..262116583e5a 100644
--- a/tests/python/topi/python/test_topi_group_conv2d.py
+++ b/tests/python/topi/python/test_topi_group_conv2d.py
@@ -52,7 +52,7 @@ def _transform_kernel(kernel, ic_bn, oc_bn):
 }
 
 _group_conv2d_nhwc_implement = {
-    "generic": (topi.nn.group_conv2d_nhwc, topi.generic.schedule_group_conv2d_nhwc),
+    "generic": (topi.nn.group_conv2d_nhwc, topi.generic.schedule_group_conv2d_nhwc)
 }
 
 
diff --git a/tests/python/topi/python/test_topi_math.py b/tests/python/topi/python/test_topi_math.py
index c7f80033bdf3..0077f42f0357 100644
--- a/tests/python/topi/python/test_topi_math.py
+++ b/tests/python/topi/python/test_topi_math.py
@@ -72,12 +72,7 @@ def check_target(target, dev):
             check_target(target, dev)
 
     def test_isnan(
-        low,
-        high,
-        shape=(20, 3),
-        dtype="float32",
-        check_round=False,
-        skip_name_check=False,
+        low, high, shape=(20, 3), dtype="float32", check_round=False, skip_name_check=False
     ):
         m = te.var("m")
         l = te.var("l")
diff --git a/tests/python/topi/python/test_topi_matmul.py b/tests/python/topi/python/test_topi_matmul.py
index de2d4d3c4c8e..02d977706318 100644
--- a/tests/python/topi/python/test_topi_matmul.py
+++ b/tests/python/topi/python/test_topi_matmul.py
@@ -46,9 +46,7 @@ def verify_nn_matmul(sa, sb, transp_a, transp_b):
     b = np.random.uniform(low=-1.0, high=1.0, size=sb).astype(np.float32)
     c1 = np.matmul(np.transpose(a) if transp_a else a, np.transpose(b) if transp_b else b)
     c2 = with_tvm(
-        lambda A, B: topi.nn.matmul(A, B, transpose_a=transp_a, transpose_b=transp_b),
-        a,
-        b,
+        lambda A, B: topi.nn.matmul(A, B, transpose_a=transp_a, transpose_b=transp_b), a, b
     )
     tvm.testing.assert_allclose(c1, c2, rtol=1e-5, atol=1e-5)
 
diff --git a/tests/python/topi/python/test_topi_relu.py b/tests/python/topi/python/test_topi_relu.py
index 83007e16f81d..7942afa5e503 100644
--- a/tests/python/topi/python/test_topi_relu.py
+++ b/tests/python/topi/python/test_topi_relu.py
@@ -30,9 +30,7 @@
 
 
 m, n, dtype = tvm.testing.parameters(
-    (10, 128, "float32"),
-    (128, 64, "float16"),
-    (1024 * 100, 512, "float32"),
+    (10, 128, "float32"), (128, 64, "float16"), (1024 * 100, 512, "float32")
 )
 
 
@@ -76,9 +74,7 @@ def test_leaky_relu(size, alpha):
 
 
 x, w, axis, weight_reshape = tvm.testing.parameters(
-    ((1, 3, 2, 2), (3,), 1, (3, 1, 1)),
-    ((1, 3, 2, 2), (2,), 2, (2, 1)),
-    ((1, 3), (3,), 1, (3,)),
+    ((1, 3, 2, 2), (3,), 1, (3, 1, 1)), ((1, 3, 2, 2), (2,), 2, (2, 1)), ((1, 3), (3,), 1, (3,))
 )
 
 
diff --git a/tests/python/topi/python/test_topi_reorg.py b/tests/python/topi/python/test_topi_reorg.py
index d2f01109e2b8..181c1ed59a86 100644
--- a/tests/python/topi/python/test_topi_reorg.py
+++ b/tests/python/topi/python/test_topi_reorg.py
@@ -23,10 +23,7 @@
 import tvm.topi.testing
 import tvm.testing
 
-_reorg_schedule = {
-    "generic": topi.generic.schedule_reorg,
-    "gpu": topi.cuda.schedule_reorg,
-}
+_reorg_schedule = {"generic": topi.generic.schedule_reorg, "gpu": topi.cuda.schedule_reorg}
 
 
 def verify_reorg(batch, in_size, in_channel, stride):
diff --git a/tests/python/topi/python/test_topi_scan.py b/tests/python/topi/python/test_topi_scan.py
index cd77a1ccfbce..017dcb768898 100644
--- a/tests/python/topi/python/test_topi_scan.py
+++ b/tests/python/topi/python/test_topi_scan.py
@@ -58,12 +58,7 @@ def get_implementations(name, axis, dtype, exclusive):
     }
 
 
-def _run_tests(
-    dev,
-    target,
-    op_name: str = "cumsum",
-    gt_func: Callable[..., np.array] = np.cumsum,
-):
+def _run_tests(dev, target, op_name: str = "cumsum", gt_func: Callable[..., np.array] = np.cumsum):
     def check_scan(np_ref, data, axis=None, dtype=None, exclusive=False):
         implementations = get_implementations(op_name, axis, dtype, exclusive)
         fcompute, fschedule = tvm.topi.testing.dispatch(target, implementations)
diff --git a/tests/python/topi/python/test_topi_sparse.py b/tests/python/topi/python/test_topi_sparse.py
index 35492c9a5f5c..446ec39c7182 100644
--- a/tests/python/topi/python/test_topi_sparse.py
+++ b/tests/python/topi/python/test_topi_sparse.py
@@ -72,8 +72,8 @@ def check_device(target, dev):
         tvm.testing.assert_allclose(d.numpy(), d_np, rtol=1e-4, atol=1e-4)
 
     # Unable to use the following idiom and thus restrict to llvm only.
-    #for target, dev in tvm.testing.enabled_targets():
-    for target in ['llvm']:
+    # for target, dev in tvm.testing.enabled_targets():
+    for target in ["llvm"]:
         dev = tvm.device(target, 0)
         check_device(target, dev)
 
@@ -114,7 +114,7 @@ def check_device(target, dev):
 
     # This test is not yet ready to use the idiom ,
     # for target, dev in tvm.testing.enabled_targets()
-    for target in ['llvm']:
+    for target in ["llvm"]:
         dev = tvm.device(target, 0)
         check_device(target, dev)
 
@@ -580,9 +580,9 @@ def check_device(target, dev):
         )
         tvm.testing.assert_allclose(Y_tvm.numpy(), Y_np.astype("float32"), atol=1e-4, rtol=1e-4)
 
-    #Unable to use the idiom below as the test isn't suitable for non llvm targets.
-    #for target, dev in tvm.testing.enabled_targets():
-    for target in ['llvm']:
+    # Unable to use the idiom below as the test isn't suitable for non llvm targets.
+    # for target, dev in tvm.testing.enabled_targets():
+    for target in ["llvm"]:
         dev = tvm.device(target, 0)
         check_device(target, dev)
 
diff --git a/tests/python/topi/python/test_topi_vision.py b/tests/python/topi/python/test_topi_vision.py
index 234107d6686e..3c2f485e5cf7 100644
--- a/tests/python/topi/python/test_topi_vision.py
+++ b/tests/python/topi/python/test_topi_vision.py
@@ -55,10 +55,7 @@
     "gpu": (topi.vision.roi_align_nchw, topi.cuda.schedule_roi_align),
 }
 
-_roi_pool_schedule = {
-    "generic": topi.generic.schedule_roi_pool,
-    "gpu": topi.cuda.schedule_roi_pool,
-}
+_roi_pool_schedule = {"generic": topi.generic.schedule_roi_pool, "gpu": topi.cuda.schedule_roi_pool}
 
 _proposal_implement = {
     "generic": (topi.vision.rcnn.proposal, topi.generic.schedule_proposal),