Fix python format

apache · Feb 11, 2021 · a1d5c43 · a1d5c43
1 parent 73330ec
commit a1d5c43
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 23 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -2300,13 +2300,12 @@ def _impl(inputs, attr, params, mod):
         assert len(inputs) == 1
         x = inputs[0]
         [output, indices, counts, num_uniq] = _op.unique(x)
-        output_sliced = _op.strided_slice(
-            output, begin=[0], end=num_uniq, slice_mode="size"
-        )
+        output_sliced = _op.strided_slice(output, begin=[0], end=num_uniq, slice_mode="size")
         return [output_sliced, indices]
 
     return _impl
 
+
 # compatible operators that do NOT require any conversion.
 _identity_list = []
 

diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py
@@ -911,4 +911,9 @@ def unique_shape_func(attrs, inputs, _):
     """
     Shape func for unique operator.
     """
-    return [_unique_shape_1(inputs[0]), _unique_shape_1(inputs[0]), _unique_shape_1(inputs[0]), _unique_shape_2(inputs[0])]
+    return [
+        _unique_shape_1(inputs[0]),
+        _unique_shape_1(inputs[0]),
+        _unique_shape_1(inputs[0]),
+        _unique_shape_2(inputs[0]),
+    ]
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
@@ -962,9 +962,7 @@ def unique(data):
     return te.extern(
         [data.shape, data.shape, data.shape, (1,)],
         [data],
-        lambda ins, outs: tvm.tir.call_packed(
-            "tvm.contrib.algorithm.unique", ins[0], *outs
-        ),
+        lambda ins, outs: tvm.tir.call_packed("tvm.contrib.algorithm.unique", ins[0], *outs),
         dtype=[data.dtype, "int32", "int32", "int32", "int32"],
         name="unique_cpu",
         tag="unique_cpu",

diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
@@ -4689,11 +4689,7 @@ def lstm_cell():
 #######################################################################
 # Unique
 # ------------
-def _test_unique(
-    n,
-    dtype,
-    is_dyn
-):
+def _test_unique(n, dtype, is_dyn):
     """ One iteration of a Stridedslice """
 
     tf.reset_default_graph()
@@ -4705,8 +4701,7 @@ def _test_unique(
             in_data = tf.constant(np_data, dtype, name="in_data")
         tf.unique(in_data)
         if is_dyn:
-            compare_tf_with_tvm(np_data, "in_data:0", [
-                                "Unique:0", "Unique:1"], mode="vm")
+            compare_tf_with_tvm(np_data, "in_data:0", ["Unique:0", "Unique:1"], mode="vm")
         else:
             compare_tf_with_tvm(None, "", ["Unique:0", "Unique:1"])
 

diff --git a/tests/python/relay/test_op_level6.py b/tests/python/relay/test_op_level6.py
@@ -142,7 +142,8 @@ def verify_topk(k, axis, ret_type, is_ascend, dtype):
 def test_unique():
     def calc_unique(data):
         uniq, index, inverse, counts = np.unique(
-            data, return_index=True, return_inverse=True, return_counts=True)
+            data, return_index=True, return_inverse=True, return_counts=True
+        )
         order = np.argsort(index)
         reverse_order = dict(zip(order, np.arange(len(order))))
         uniq = uniq[order].astype(data.dtype)
@@ -168,21 +169,18 @@ def verify_unique(n, dtype, is_dyn=False):
         for target, ctx in tvm.testing.enabled_targets():
             for kind in backends:
                 mod = tvm.ir.IRModule.from_expr(func)
-                intrp = relay.create_executor(
-                    kind, mod=mod, ctx=ctx, target=target)
+                intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target)
                 op_res = intrp.evaluate()(x_data)
                 ref_res = calc_unique(x_data)
                 num_uniq = ref_res[3][0]
                 assert num_uniq == op_res[3].asnumpy()[0]
                 # output
-                tvm.testing.assert_allclose(
-                    op_res[0].asnumpy()[:num_uniq], ref_res[0], rtol=1e-5)
+                tvm.testing.assert_allclose(op_res[0].asnumpy()[:num_uniq], ref_res[0], rtol=1e-5)
                 # inverse_indices
-                tvm.testing.assert_allclose(
-                    op_res[1].asnumpy(), ref_res[1], rtol=1e-5)
+                tvm.testing.assert_allclose(op_res[1].asnumpy(), ref_res[1], rtol=1e-5)
                 # count
-                tvm.testing.assert_allclose(
-                    op_res[2].asnumpy()[:num_uniq], ref_res[2], rtol=1e-5)
+                tvm.testing.assert_allclose(op_res[2].asnumpy()[:num_uniq], ref_res[2], rtol=1e-5)
+
     for dtype in ["int32", "int64"]:
         for is_dyn in [False, True]:
             verify_unique((50), dtype, is_dyn=is_dyn)