pylint + fix comments

Signed-off-by: xavier dupré <[email protected]>

Author: sdpython

tests/test_backend.py

@@ -24,7 +24,7 @@
 from tf2onnx import constants, utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tf_loader import is_tf2
-from tf2onnx.onnx_opset.signal import DFT_constant
+from tf2onnx.onnx_opset.signal import make_dft_constant
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,function-redefined,cell-var-from-loop
 
@@ -3633,18 +3633,17 @@ def DFT_slow(x, M):
             return np.transpose(res, (0, 2, 1))
 
         x_val = make_xval([2, 4]).astype(np.float32)
-        M_both = DFT_constant(x_val.shape[1], x_val.dtype, x_val.shape[1])
+        M_both = make_dft_constant(x_val.shape[1], x_val.dtype, x_val.shape[1])
         fft = DFT_slow(x_val, M_both)
         fft_npy = np.fft.rfft(x_val)
         assert_almost_equal(fft[0, :, :], np.real(fft_npy))
         assert_almost_equal(fft[1, :, :], np.imag(fft_npy))
 
-        for op in [tf.signal.rfft]:
-            x_val = make_xval([3, 4]).astype(np.float32)
-            def func(x):
-                op_ = op(x)
-                return tf.abs(op_, name=_TFOUTPUT)
-            self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+        x_val = make_xval([3, 4]).astype(np.float32)
+        def func(x):
+            op_ = tf.signal.rfft(x)
+            return tf.abs(op_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
 
 if __name__ == '__main__':

tf2onnx/onnx_opset/signal.py

@@ -14,23 +14,22 @@
 import numpy as np
 from onnx import onnx_pb
 from onnx.numpy_helper import to_array
-from tf2onnx import constants, utils
+from tf2onnx import utils
 from tf2onnx.handler import tf_op
-from tf2onnx.onnx_opset import common
 
 logger = logging.getLogger(__name__)
 
 
 # pylint: disable=unused-argument,missing-docstring
 
-def DFT_constant(N, dtype, fft_length):
+def make_dft_constant(N, dtype, fft_length):
     n = np.arange(N)
     k = n.reshape((N, 1)).astype(np.float64)
-    M = np.exp(-2j * np.pi * k * n / N)
-    M = M[:fft_length // 2 + 1]
-    both = np.empty((2, ) + M.shape, dtype=dtype)
-    both[0, :, :] = np.real(M)
-    both[1, :, :] = np.imag(M)
+    mat = np.exp(-2j * np.pi * k * n / N)
+    mat = mat[:fft_length // 2 + 1]
+    both = np.empty((2,) + mat.shape, dtype=dtype)
+    both[0, :, :] = np.real(mat)
+    both[1, :, :] = np.imag(mat)
     return both
 
 
@@ -52,7 +51,7 @@ def version_1(cls, ctx, node, **kwargs):
         <https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/>`_.
 
         Complex version:
-        
+
         ::
 
             import numpy as np
@@ -71,7 +70,7 @@ def DFT(x, fft_length=None):
                 if fft_length is None:
                     fft_length = x.shape[0]
                 cst = _DFT_cst(x.shape[0], fft_length)
-                return np.dot(cst, x).T                
+                return np.dot(cst, x).T
 
         Real version, first axis is (real, imag) part:
 
@@ -84,7 +83,7 @@ def _DFT_real_cst(N, fft_length):
                 k = n.reshape((N, 1)).astype(np.float64)
                 M = np.exp(-2j * np.pi * k * n / N)
                 M = M[:fft_length // 2 + 1]
-                both = np.empty((2, ) + M.shape)
+                both = np.empty((2,) + M.shape)
                 both[0, :, :] = np.real(M)
                 both[1, :, :] = np.imag(M)
                 return both
@@ -104,7 +103,7 @@ def DFT_real(x, fft_length=None):
         onnx_dtype = ctx.get_dtype(node.input[0])
         shape = ctx.get_shape(node.input[0])
         np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
-        N = shape[-1]
+        shape_n = shape[-1]
         utils.make_sure(len(node.input) == 2, "Two inputs expected not %r", len(node.input))
 
         # This input should be a constant.
@@ -114,13 +113,13 @@ def DFT_real(x, fft_length=None):
                         "fft_length should be a constant, the other case is not implemented yet.")
         value = node_fft_length.get_attr("value")
         value_array = to_array(value.t)
-        utils.make_sure(value_array.shape == (1, ), "Unexpected shape for fft_length (%r)", value_array.shape)
+        utils.make_sure(value_array.shape == (1,), "Unexpected shape for fft_length (%r)", value_array.shape)
         fft_length = value_array[0]
 
         # TODO: handle this parameter when onnx.helper.make_node is fixed.
         # Tcomplex = node.get_attr("Tcomplex")
 
-        if np_dtype in (np.float16, ):
+        if np_dtype == np.float16:
             res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float16)
             np_dtype = np.float16
         elif np_dtype in (np.float32, np.complex64):
@@ -130,9 +129,9 @@ def DFT_real(x, fft_length=None):
             res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float64)
             np_dtype = np.float64
 
-        real_imag_part = DFT_constant(N, np_dtype, fft_length)
+        real_imag_part = make_dft_constant(shape_n, np_dtype, fft_length)
         onx_real_imag_part = ctx.make_const(
-            name=utils.make_name('cst_rfft_%d' % N), np_val=real_imag_part)
+            name=utils.make_name('cst_rfft_%d' % shape_n), np_val=real_imag_part)
 
         shapei = list(np.arange(len(shape)))
         perm = shapei[:-2] + [shapei[-1], shapei[-2]]
@@ -150,7 +149,7 @@ def DFT_real(x, fft_length=None):
         perm = shapei[:-2] + [shapei[-1], shapei[-2]]
         last_node = ctx.make_node(
             "Transpose", inputs=[mult.output[0]], attr=dict(perm=perm),
-            name=utils.make_name('CPLX_' + node.name + 'rfft'),            
+            name=utils.make_name('CPLX_' + node.name + 'rfft'),
             shapes=[new_shape], dtypes=[res_onnx_dtype])
 
         ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()
@@ -225,5 +224,5 @@ def version_1(cls, ctx, node, **kwargs):
         cls.any_version(1, ctx, node, **kwargs)
 
     @classmethod
-    def version_11(cls, ctx, node, **kwargs):
+    def version_13(cls, ctx, node, **kwargs):
         cls.any_version(11, ctx, node, **kwargs)

tf2onnx/tf_utils.py

@@ -236,8 +236,8 @@ def tflist_to_onnx(g, shape_override, const_node_values=None):
                     "key_dtype", "value_dtype", "Tin", "Tout", "capacity", "component_types", "shapes",
                     "Toutput_types",
                     "Tcomplex", "Treal",  # For RFFT, Tcomplex is ignored because
-                                         # onnx.helper.make_node fails,
-                                         # TODO: it should be added back.
+                                          # onnx.helper.make_node fails,
+                                          # TODO: it should be added back.
                     }
 
     node_list = g.get_operations()