version_converter/adapters/split_17_18.h

0001 // Copyright (c) ONNX Project Contributors
0002
0003 /*
0004  * SPDX-License-Identifier: Apache-2.0
0005  */
0006
0007 // Adapter for Split in default domain from version 17 to 18
0008
0009 #pragma once
0010
0011 #include <memory>
0012
0013 #include "onnx/version_converter/adapters/adapter.h"
0014 #include "onnx/version_converter/adapters/transformers.h"
0015
0016 namespace ONNX_NAMESPACE {
0017 namespace version_conversion {
0018
0019 class Split_17_18 : public Adapter {
0020  public:
0021   explicit Split_17_18() : Adapter("Split", OpSetID(17), OpSetID(18)) {}
0022
0023   void adapt_split_17_18(std::shared_ptr<Graph>, Node* node) const {
0024     const auto num_outputs = node->outputs().size();
0025     SetAttribute(knum_outputs, num_outputs);
0026   }
0027
0028   Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
0029     // if node does not have neither 'num_outputs' attribute nor 'split' input
0030     if (!node->hasAttribute(knum_outputs) && node->inputs().size() != 2) {
0031       adapt_split_17_18(graph, node);
0032     }
0033     return node;
0034   }
0035 };
0036
0037 } // namespace version_conversion
0038 } // namespace ONNX_NAMESPACE