Merge pull request #1116 from xadupre/rttf2

Better error message when it fails due to RFFT

Author: xadupre

tests/test_backend.py

@@ -3659,10 +3659,22 @@ def dft_slow(x, M):
         assert_almost_equal(fft[1, :, :], np.imag(fft_npy))
 
         x_val = make_xval([3, 4]).astype(np.float32)
-        def func(x):
+        def func1(x):
             op_ = tf.signal.rfft(x)
             return tf.abs(op_, name=_TFOUTPUT)
-        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+        self._run_test_case(func1, [_OUTPUT], {_INPUT: x_val})
+
+        def func2(x):
+            op_ = tf.signal.rfft(x)
+            return tf.cos(op_, name=_TFOUTPUT)
+        with self.assertRaises(ValueError):
+            self._run_test_case(func2, [_OUTPUT], {_INPUT: x_val})
+
+        def func3(x):
+            op_ = tf.signal.rfft(x)
+            return tf.identity(op_, name=_TFOUTPUT)
+        with self.assertRaises(ValueError):
+            self._run_test_case(func3, [_OUTPUT], {_INPUT: x_val})
 
 
 if __name__ == '__main__':

tf2onnx/onnx_opset/signal.py

@@ -99,6 +99,12 @@ def DFT_real(x, fft_length=None):
                 res = np.dot(cst, x)
                 return np.transpose(res, (0, 2, 1))
         """
+        consumers = ctx.find_output_consumers(node.output[0])
+        consumer_types = set(op.type for op in consumers)
+        utils.make_sure(
+            consumer_types == {'ComplexAbs'},
+            "Current implementation of RFFT only allows ComplexAbs as consumer not %r",
+            consumer_types)
 
         onnx_dtype = ctx.get_dtype(node.input[0])
         shape = ctx.get_shape(node.input[0])