[BugFix] add the default value for DFT in ONNX frontend (#16659)

apache · Mar 4, 2024 · 46aaf61 · 46aaf61
1 parent 880af30
commit 46aaf61
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
@@ -4809,9 +4809,9 @@ class DFT(OnnxOpConverter):
     @classmethod
     def _impl_v17(cls, inputs, attr, params):
         # ************************* Read attrs *************************
-        axis = attr.get("axis")
-        inverse = attr.get("inverse")
-        onesided = attr.get("onesided")
+        axis = attr.get("axis", 1)
+        inverse = attr.get("inverse", 0)
+        onesided = attr.get("onesided", 0)
 
         # ************************* Read inputs ************************
         input_tensor = inputs[0]

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
@@ -8238,7 +8238,7 @@ def verify_dft(
     D = 7
 
     for axis in list(range(1, n)) + [-2]:
-        for inverse, onesided in [(0, 0), (0, 1), (1, 0)]:
+        for inverse, onesided in [(0, 0), (0, 1), (1, 0), (None, None)]:
             for n_fft in [D, D - 1, D + 1]:
                 for c in [1, 2]:
                     input_shape = [batch_size] + n * [D] + [c]