diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 0de00517e9a8a..50221c7baf288 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -19,6 +19,7 @@ """TF: Tensorflow frontend.""" import warnings from collections import defaultdict +from collections import deque # Numpy support import numpy as np @@ -1778,7 +1779,6 @@ def _impl(inputs, attr, params, mod): s0 = list(s0.asnumpy().reshape([-1])) s1 = list(s1.asnumpy().reshape([-1])) s0_size, s1_size = len(s0), len(s1) - from collections import deque out = deque([]) for i in range(1, min(s0_size, s1_size) + 1): diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 2e4237b34ab14..56e3799850138 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1725,9 +1725,7 @@ def test_read_variable_op(target, dev): # Now convert the variables to constant and run inference on the converted graph final_graph_def = tf.graph_util.convert_variables_to_constants( - sess, - sess.graph.as_graph_def(add_shapes=True), - out_node, + sess, sess.graph.as_graph_def(add_shapes=True), out_node, ) tvm_output = run_tvm_graph( @@ -2228,20 +2226,10 @@ def _test_sparse_segment_variant( ) @pytest.mark.parametrize("use_dyn", [True, False]) @pytest.mark.parametrize( - "tf_op", - [ - tf.sparse.segment_sum, - tf.sparse.segment_sqrt_n, - tf.sparse.segment_mean, - ], + "tf_op", [tf.sparse.segment_sum, tf.sparse.segment_sqrt_n, tf.sparse.segment_mean,], ) def test_forward_sparse_segment_sum_variants( - tf_op, - data_np, - indices_np, - segment_ids_np, - num_segments, - use_dyn, + tf_op, data_np, indices_np, segment_ids_np, num_segments, use_dyn, ): """sparse segment sum variants tests""" _test_sparse_segment_variant(tf_op, data_np, indices_np, segment_ids_np, num_segments, use_dyn) @@ -2269,10 +2257,7 @@ def _test_math_segment_sum(data_np, segment_ids_np, use_dyn=False): _ = tf.math.segment_sum(data, segment_ids, name="segment_sum") compare_tf_with_tvm( - [data_np, segment_ids_np], - [data.name, segment_ids.name], - ["segment_sum:0"], - mode="vm", + [data_np, segment_ids_np], [data.name, segment_ids.name], ["segment_sum:0"], mode="vm", ) @@ -2287,18 +2272,12 @@ def _test_math_segment_sum(data_np, segment_ids_np, use_dyn=False): np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64), np.array([0, 0, 1], dtype=np.int32), ), - ( - np.random.random((6, 4, 5)), - np.array([0, 0, 1, 2, 2, 3], dtype=np.int64), - ), + (np.random.random((6, 4, 5)), np.array([0, 0, 1, 2, 2, 3], dtype=np.int64),), ( np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32), np.array([0, 0, 1], dtype=np.int32), ), - ( - np.random.random((9, 4, 5, 7)), - np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int64), - ), + (np.random.random((9, 4, 5, 7)), np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int64),), ], ) @pytest.mark.parametrize("use_dyn", [True, False]) @@ -2961,9 +2940,7 @@ def test_forward_multi_output(): with tf.Session() as sess: final_graph_def = tf.graph_util.convert_variables_to_constants( - sess, - sess.graph.as_graph_def(add_shapes=True), - out_node, + sess, sess.graph.as_graph_def(add_shapes=True), out_node, ) tf_output = run_tf_graph(sess, in_data, in_name, out_name) tvm_output = run_tvm_graph( @@ -3438,11 +3415,7 @@ def _get_tensorflow_output(): sess.run([variables.global_variables_initializer()]) res = sess.run( [g, out_m0, out_m1], - { - x.name: np.array([[1.0, 1.0]]), - m0.name: in_state_c, - m1.name: in_state_h, - }, + {x.name: np.array([[1.0, 1.0]]), m0.name: in_state_c, m1.name: in_state_h,}, ) graph_def = sess.graph.as_graph_def(add_shapes=True) final_graph_def = graph_util.convert_variables_to_constants( @@ -3620,7 +3593,7 @@ def test_forward_logical(): ####################################################################### -# Where, Select +# Where, Select, SelectV2 # ------------- def test_forward_where(): """ Where: return elements depending on conditions""" @@ -4843,9 +4816,7 @@ def _test_identityn(data_np_list): output = tf.identity_n(data_tensors) output_names = [out.name for out in output] compare_tf_with_tvm( - data_np_list, - data_tensors_name, - output_names, + data_np_list, data_tensors_name, output_names, ) @@ -4867,12 +4838,7 @@ def _test_identityn(data_np_list): np.array([True, False, True]), ] ), - ( - [ - np.array([]), - np.array([[]]), - ] - ), + ([np.array([]), np.array([[]]),]), ], ) def test_forward_identityn(data_np_list):