DFT-20 version converter

onnx/test/version_converter/automatic_conversion_test_base.py

@@ -1,6 +1,8 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
 import string
 import unittest
 from typing import Any, Dict, List, Optional, Sequence, Union, cast
@@ -12,6 +14,29 @@
 
 
 class TestAutomaticConversion(unittest.TestCase):
+    def _test_model_conversion(
+        self, to_opset: int, model: str | onnx.ModelProto
+    ) -> None:
+        if isinstance(model, str):
+            model = onnx.parser.parse_model(model)
+        onnx.checker.check_model(model)
+        shape_inference.infer_shapes(model, strict_mode=True)
+
+        converted = version_converter.convert_version(model, to_opset)
+        onnx.checker.check_model(converted)
+        shape_inference.infer_shapes(converted, strict_mode=True)
+
+    def _test_model_conversion_fails(
+        self, to_opset: int, model: str | onnx.ModelProto
+    ) -> None:
+        if isinstance(model, str):
+            model = onnx.parser.parse_model(model)
+        onnx.checker.check_model(model)
+        shape_inference.infer_shapes(model, strict_mode=True)
+
+        with self.assertRaises(RuntimeError):
+            version_converter.convert_version(model, to_opset)
+
     def _test_op_conversion(
         self,
         op: str,
@@ -117,9 +142,4 @@ def _test_op_conversion(
             producer_name="test",
             opset_imports=[helper.make_opsetid("", start_opset)],
         )
-        onnx.checker.check_model(original)
-        shape_inference.infer_shapes(original, strict_mode=True)
-
-        converted = version_converter.convert_version(original, end_opset)
-        onnx.checker.check_model(converted)
-        shape_inference.infer_shapes(converted, strict_mode=True)
+        self._test_model_conversion(end_opset, original)

onnx/test/version_converter/automatic_downgrade_test.py

@@ -49,6 +49,56 @@ def test_reduce_ops(self, op) -> None:
             initializer=[axes],
         )
 
+    def test_dft20_no_axis(self) -> None:
+        self._test_model_conversion(
+            to_opset=19,
+            model="""
+            <ir_version: 9, opset_import: [ "" : 20]>
+            dft_no_axis (float[N, M, 1] x) => (float[N, M, 2] y)
+            {
+                y = DFT (x)
+            }
+        """,
+        )
+
+    def test_dft20_initializer_axis(self) -> None:
+        self._test_model_conversion(
+            to_opset=19,
+            model="""
+            <ir_version: 9, opset_import: [ "" : 20]>
+            dft_no_axis (float[N, M, 1] x, int64 dft_length) => (float[N, K, 2] y)
+            <int64 axis = {1}>
+            {
+                y = DFT (x, dft_length, axis)
+            }
+        """,
+        )
+
+    def test_dft20_constant_axis(self) -> None:
+        self._test_model_conversion(
+            to_opset=19,
+            model="""
+            <ir_version: 9, opset_import: [ "" : 20]>
+            dft_no_axis (float[N, M, 1] x, int64 dft_length) => (float[N, K, 2] y)
+            {
+                axis = Constant <value = int64{1}>()
+                y = DFT (x, dft_length, axis)
+            }
+        """,
+        )
+
+    def test_dft20_unknown_axis(self) -> None:
+        self._test_model_conversion_fails(
+            to_opset=19,
+            model="""
+            <ir_version: 9, opset_import: [ "" : 20]>
+            dft_no_axis (float[N, M, 1] x, int64 dft_length, int64 axis) => (float[P, K, 2] y)
+            {
+                y = DFT (x, dft_length, axis)
+            }
+        """,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

onnx/test/version_converter/automatic_upgrade_test.py

@@ -5,7 +5,6 @@
 
 import automatic_conversion_test_base
 import numpy as np
-import pytest
 
 import onnx
 from onnx import TensorProto, helper
@@ -1442,7 +1441,6 @@ def test_HannWindow(self) -> None:
     def test_HammingWindow(self) -> None:
         self._test_window_function("HammingWindow")
 
-    @pytest.mark.xfail(reason="FIXME(#5613): Implement version converters for DFT")
     def test_DFT(self) -> None:
         self._test_op_upgrade("DFT", 17, [[2, 16, 1], []], [[2, 16, 2]])
         self._test_op_upgrade("DFT", 17, [[2, 16, 2], []], [[2, 16, 2]])
@@ -1458,6 +1456,9 @@ def test_DFT(self) -> None:
         self._test_op_upgrade(
             "DFT", 17, [[2, 16, 2], []], [[2, 16, 2]], attrs={"inverse": 1}
         )
+        self._test_op_upgrade(
+            "DFT", 17, [[2, 16, 2], []], [[2, 16, 2]], attrs={"inverse": 1, "axis": 0}
+        )
 
     def _test_short_time_fourier_transform(self, operator_name: str) -> None:
         # Real

onnx/version_converter/adapters/axis_attribute_to_input.h

@@ -0,0 +1,74 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+
+class AxisAttributeToInput : public Adapter {
+ public:
+  AxisAttributeToInput(
+      const std::string& op_name,
+      const OpSetID& initial,
+      const OpSetID& target,
+      size_t axis_index,
+      int64_t default_axis)
+      : Adapter(op_name, initial, target), axis_index(axis_index), default_axis(default_axis) {}
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    if (node->hasAttribute(kaxis)) {
+      AttrToInput(graph, node, node->i(kaxis), this->axis_index);
+      node->removeAttribute(kaxis);
+      return node;
+    }
+
+    // Fill in the default value for axis
+    AttrToInput(graph, node, default_axis, this->axis_index);
+    return node;
+  }
+
+ private:
+  size_t axis_index;
+  int64_t default_axis;
+
+  void AttrToInput(std::shared_ptr<Graph> graph, Node* node, int64_t axis, size_t axis_index) const {
+    const ArrayRef<Value*>& inputs = node->inputs();
+
+    // Add the optional inputs if they don't exist
+    for (size_t i = inputs.size(); i < axis_index; ++i) {
+      Node* empty_input = graph->create(kUndefined);
+      empty_input->insertBefore(node);
+      node->addInput(empty_input->output());
+    }
+
+    // Add the axis input
+    Node* constant = CreateAxisInput(graph, node, axis);
+    node->addInput(constant->output());
+  }
+
+  Node* CreateAxisInput(std::shared_ptr<Graph> graph, Node* node, int64_t axis) const {
+    Tensor t;
+    t.elem_type() = TensorProto_DataType_INT64;
+    t.sizes() = std::vector<int64_t>{};
+    auto& data = t.int64s();
+    data.emplace_back(axis);
+
+    Node* constant = graph->create(kConstant);
+    constant->insertBefore(node);
+    constant->t_(kvalue, t);
+    return constant;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

onnx/version_converter/adapters/axis_input_to_attribute.h

@@ -0,0 +1,99 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "onnx/version_converter/adapters/adapter.h"
+
+namespace ONNX_NAMESPACE {
+namespace version_conversion {
+class AxisInputToAttribute : public Adapter {
+ public:
+  explicit AxisInputToAttribute(
+      const std::string& op_name,
+      const OpSetID& initial,
+      const OpSetID& target,
+      size_t axis_index,
+      int64_t default_axis)
+      : Adapter(op_name, initial, target), axis_index(axis_index), default_axis(default_axis) {}
+
+  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
+    if (!HasAxisInput(node)) {
+      node->i_(kaxis, this->default_axis);
+      return EnsureAndReturnNode(node);
+    }
+
+    const ArrayRef<Value*>& inputs = node->inputs();
+    Value* axis_val = inputs[this->axis_index];
+    Node* axis_node = axis_val->node();
+
+    if (axis_node->kind() == kConstant) {
+      HandleConstantNode(node, axis_node, axis_val);
+      return EnsureAndReturnNode(node);
+    }
+
+    if (graph->is_constant_initializer(axis_val)) {
+      HandleInitializerNode(graph, node, axis_val);
+      return EnsureAndReturnNode(node);
+    }
+
+    ONNX_ASSERTM(false, "Axis input must be a constant or initializer for promotion to attribute.");
+  }
+
+ private:
+  size_t axis_index;
+  int64_t default_axis;
+
+  bool HasAxisInput(const Node* node) const {
+    const ArrayRef<const Value*>& inputs = node->inputs();
+    return inputs.size() > this->axis_index && inputs[this->axis_index]->node()->kind() != kUndefined;
+  }
+
+  void HandleConstantNode(Node* node, Node* axis_node, Value* axis_val) const {
+    const std::vector<int64_t>& int64s = axis_node->t(kvalue).int64s();
+    if (int64s.empty()) {
+      std::string raw_data = axis_node->t(kvalue).raw();
+      ONNX_ASSERTM(
+          raw_data.size() != 0 && raw_data.size() % 8 == 0,
+          "Raw Data must be non-empty and size must be a multiple of 8");
+      const int64_t* raw = reinterpret_cast<const int64_t*>(raw_data.c_str());
+      node->i_(kaxis, raw[0]);
+    } else {
+      node->i_(kaxis, int64s.at(0));
+    }
+    node->removeInput(this->axis_index);
+    if (axis_val->uses().size() < 1) {
+      axis_node->destroy();
+    }
+  }
+
+  void HandleInitializerNode(std::shared_ptr<Graph> graph, Node* node, Value* axis_val) const {
+    const std::string initializer_name = axis_val->uniqueName();
+    for (const auto& initializer : graph->initializers()) {
+      if (initializer.name() == initializer_name) {
+        node->i_(kaxis, initializer.int64s().at(0));
+        node->removeInput(this->axis_index);
+        // Remove initializer
+        if (axis_val->uses().size() < 1)
+          graph->eraseInitializer(initializer_name);
+        break;
+      }
+    }
+  }
+
+  inline Node* EnsureAndReturnNode(Node* node) const {
+    ONNX_ASSERTM(node->hasAttribute(kaxis), "Axis attribute not created. This may be a bug.");
+    return node;
+  }
+};
+
+} // namespace version_conversion
+} // namespace ONNX_NAMESPACE

onnx/version_converter/convert.h

@@ -19,6 +19,8 @@
 #include "onnx/version_converter/BaseConverter.h"
 #include "onnx/version_converter/adapters/axes_attribute_to_input.h"
 #include "onnx/version_converter/adapters/axes_input_to_attribute.h"
+#include "onnx/version_converter/adapters/axis_attribute_to_input.h"
+#include "onnx/version_converter/adapters/axis_input_to_attribute.h"
 #include "onnx/version_converter/adapters/batch_normalization_13_14.h"
 #include "onnx/version_converter/adapters/broadcast_backward_compatibility.h"
 #include "onnx/version_converter/adapters/broadcast_forward_compatibility.h"
@@ -575,9 +577,10 @@ class DefaultVersionConverter : public BaseVersionConverter {
     registerAdapter(std::make_unique<CompatibleAdapter>("Size", OpSetID(18), OpSetID(19)));
 
     /******** 19 -> 20 ********/
+    registerAdapter(std::make_unique<AxisAttributeToInput>("DFT", OpSetID(19), OpSetID(20), 2, 1));
     registerAdapter(std::make_unique<CompatibleAdapter>("ConstantOfShape", OpSetID(19), OpSetID(20)));
-    registerAdapter(std::make_unique<CompatibleAdapter>("IsNaN", OpSetID(19), OpSetID(20)));
     registerAdapter(std::make_unique<CompatibleAdapter>("IsInf", OpSetID(19), OpSetID(20)));
+    registerAdapter(std::make_unique<CompatibleAdapter>("IsNaN", OpSetID(19), OpSetID(20)));
     registerAdapter(std::make_unique<CompatibleAdapter>("ReduceMax", OpSetID(19), OpSetID(20)));
     registerAdapter(std::make_unique<CompatibleAdapter>("ReduceMin", OpSetID(19), OpSetID(20)));
     registerAdapter(std::make_unique<GridSample_19_20>());
@@ -597,6 +600,7 @@ class DefaultVersionConverter : public BaseVersionConverter {
         TensorProto_DataType_FLOAT8E5M2,
         TensorProto_DataType_FLOAT8E5M2FNUZ};
     registerAdapter(std::make_unique<TypeRestriction>("IsInf", OpSetID(19), OpSetID(20), is_inf_10_unallowed_types));
+    registerAdapter(std::make_unique<AxisInputToAttribute>("DFT", OpSetID(20), OpSetID(19), 2, -2));
     const std::vector<TensorProto_DataType> reduce_min_max_18_unallowed_types = {TensorProto_DataType_BOOL};
     registerAdapter(
         std::make_unique<TypeRestriction>("ReduceMax", OpSetID(20), OpSetID(19), reduce_min_max_18_unallowed_types));