Update RFFT2D to handle 3D arrays (Merge pull request #1649 from xadupre/tflfft)

Update RFFT2D to handle 3D arrays

Author: xadupre

tests/test_backend.py

@@ -5392,7 +5392,6 @@ def func(input_val):
             self.config.opset = current_opset
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
     def test_rfft_ops(self):
 
         def dft_slow(x, M, fft_length):
@@ -5434,7 +5433,6 @@ def func3(x):
             self._run_test_case(func3, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
     @skip_tfjs("TFJS executes rfft with poor accuracy")
     @check_opset_min_version(10, "Slice")
     def test_rfft_ops_fft_length(self):
@@ -5446,7 +5444,6 @@ def func1_length(x):
         self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
     @skip_tfjs("TFJS executes rfft with poor accuracy")
     @check_opset_min_version(10, "Slice")
     def test_rfft_ops_fft_length_many(self):
@@ -5461,7 +5458,6 @@ def func1_length(x):
                         self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
     @check_opset_min_version(10, "Slice")
     def test_rfft_ops_fft_length_many_bigger(self):
         for i in range(4, 7):
@@ -5491,8 +5487,7 @@ def func1_length(x):
                         self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
-    @check_opset_min_version(10, "Slice")
+    @check_opset_min_version(11, "CumSum")
     def test_rfft2d_ops(self):
 
         x_val = make_xval([3, 4]).astype(np.float32)
@@ -5516,34 +5511,35 @@ def func3(x):
             self._run_test_case(func3, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
-    @check_opset_min_version(10, "Slice")
+    @check_opset_min_version(11, "CumSum")
     def test_rfft2d_ops_fft_length(self):
 
         x_val = make_xval([3, 4]).astype(np.float32)
         def func1_length(x):
             op_ = tf.signal.rfft2d(x, np.array([3, 3], dtype=np.int32))
             return tf.abs(op_, name=_TFOUTPUT)
-        self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val}, optimize=False)
-        self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
+        with self.subTest(optimize=False):
+            self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val}, optimize=False)
+        with self.subTest(optimize=True):
+            self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
 
     @check_tf_min_version("1.14")
-    #@skip_tflite("FlexRFFT2D")
-    @check_opset_min_version(10, "Slice")
+    @check_opset_min_version(11, "CumSum")
     def test_rfft2d_ops_fft_length_many(self):
         for i in range(7, 4, -1):
             for j in range(7, 4, -1):
                 for m in range(0, 3):
                     for n in range(0, 3):
-                        with self.subTest(shape=(i, j), fft_length=(m, n)):
-                            x_val = make_xval([i, j]).astype(np.float32) / 100
-                            def func1_length(x):
-                                op_ = tf.signal.rfft2d(x, np.array([i-m, j-n], dtype=np.int32))
-                                return tf.abs(op_, name=_TFOUTPUT)
-                            self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val})
+                        for opt in [False, True]:
+                            with self.subTest(shape=(i, j), fft_length=(m, n), optimize=opt):
+                                x_val = make_xval([i, j]).astype(np.float32) / 100
+                                def func1_length(x):
+                                    op_ = tf.signal.rfft2d(x, np.array([i-m, j-n], dtype=np.int32))
+                                    return tf.abs(op_, name=_TFOUTPUT)
+                                self._run_test_case(func1_length, [_OUTPUT], {_INPUT: x_val}, optimize=opt)
 
     @check_tf_min_version("1.14")
-    @check_opset_min_version(10, "Slice")
+    @check_opset_min_version(11, "CumSum")
     @unittest.skipIf(True, reason="Not fully implemented for dynamic shape.")
     def test_fft_ops(self):
         x_val = make_xval([3, 4]).astype(np.float32)
@@ -5566,6 +5562,37 @@ def func(x):
         x_val = np.array([1, 5, 2, 0, 3, 4], dtype=np.int64)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_tf_min_version("1.14")
+    @check_opset_min_version(11, "CumSum")
+    def test_rfft2d_ops_specific_dimension(self):
+
+        x_val = make_xval([3, 1, 4]).astype(np.float32)
+
+        def func1(x):
+            op_ = tf.signal.rfft2d(x, np.array([1, 4], dtype=np.int32))
+            return tf.abs(op_, name=_TFOUTPUT)
+        with self.subTest(shape=(3, 1, 4), fft_length=(1, 4), optimize=False):
+            self._run_test_case(func1, [_OUTPUT], {_INPUT: x_val}, optimize=False)
+        with self.subTest(shape=(3, 1, 4), fft_length=(1, 4), optimize=True):
+            self._run_test_case(func1, [_OUTPUT], {_INPUT: x_val})
+
+        for shape in [(3, 1, 4), (5, 7), (3, 5, 7), (7, 5)]:
+            for fft_length in [shape[-2:], (1, shape[-1]),
+                               (min(2, shape[-2]), shape[-1]),
+                               (shape[-2], 2),
+                               (min(3, shape[-2]), min(4, shape[-2]))]:
+                if fft_length == (1, 1):
+                    # The code fails in this case but that's unlikely to happen.
+                    continue
+                for optimize in [False, True]:
+                    with self.subTest(shape=shape, fft_length=fft_length, optimize=optimize):
+                        x_val = make_xval(list(shape)).astype(np.float32)
+                        x_val /= x_val.size
+                        def func1(x):
+                            op_ = tf.signal.rfft2d(x, np.array(fft_length, dtype=np.int32))
+                            return tf.abs(op_, name=_TFOUTPUT)
+                        self._run_test_case(func1, [_OUTPUT], {_INPUT: x_val}, optimize=optimize)
+
     @check_tf_min_version("2.1")
     @skip_tflite("TFlite errors on some attributes")
     @check_opset_min_version(9, "string")

tests/tfhub/_tools.py

@@ -403,9 +403,15 @@ def call_tflite(inp):
         with open(onnx_name, "rb") as f:
             model_onnx = onnx.load(f)
 
-        call_tflite(imgs[0])
+        interpreter_details = tf.lite.Interpreter(tname, experimental_preserve_all_tensors=True)
+        input_details = interpreter_details.get_input_details()
+        index_in = input_details[0]['index']
+        interpreter_details.allocate_tensors()
+        interpreter_details.set_tensor(index_in, imgs[0])
+        interpreter_details.invoke()
+        details = interpreter_details.get_tensor_details()
+
         inputs = {input_name: imgs[0]}
-        details = interpreter.get_tensor_details()
         names_index = {}
         for tt in details:
             names_index[tt['name']] = (tt['index'], tt['quantization'], tt['quantization_parameters'])
@@ -414,7 +420,7 @@ def call_tflite(inp):
         for name_tfl, name_ort in names:
             index = names_index[name_tfl]
 
-            tfl_value = interpreter.get_tensor(index[0])
+            tfl_value = interpreter_details.get_tensor(index[0])
 
             new_name = onnx_name + ".%s.onnx" % name_ort.replace(":", "_").replace(";", "_").replace("/", "_")
             if not os.path.exists(new_name):