version_converter/adapters/split_13_12.h

0001 // Copyright (c) ONNX Project Contributors
0002
0003 /*
0004  * SPDX-License-Identifier: Apache-2.0
0005  */
0006
0007 // Adapter for all ops that remove consumed_inputs
0008
0009 #pragma once
0010
0011 #include <memory>
0012 #include <string>
0013 #include <utility>
0014 #include <vector>
0015
0016 #include "onnx/version_converter/adapters/adapter.h"
0017
0018 namespace ONNX_NAMESPACE {
0019 namespace version_conversion {
0020
0021 class Split_13_12 : public Adapter {
0022  public:
0023   explicit Split_13_12() : Adapter("Split", OpSetID(13), OpSetID(12)) {}
0024
0025   Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
0026     // Identify if 'split' is statically determined; if so, feed as attribute
0027     const ArrayRef<Value*>& inputs = node->inputs();
0028     // Get 'split' from initializer or constant operator
0029     // Identify whether we have a Constant Op or an Initializer
0030     Value* const_val = inputs[1];
0031     Node* node_ptr = const_val->node();
0032     if (node_ptr->kind() == kConstant) {
0033       // Get value attribute of kConstant
0034       const std::vector<int64_t>& int64s = node_ptr->t(kvalue).int64s();
0035       if (int64s.empty()) {
0036         // Also handle raw data
0037         std::string raw_data = node_ptr->t(kvalue).raw();
0038         ONNX_ASSERTM(
0039             raw_data.size() != 0 && raw_data.size() % 8 == 0,
0040             "Raw Data must be non-empty and size must be a multiple of 8");
0041         int64_t* raw = (int64_t*)const_cast<char*>(raw_data.c_str());
0042         node->is_(ksplit, std::vector<int64_t>(raw, raw + node_ptr->t(kvalue).size_from_dim(0)));
0043       } else {
0044         node->is_(ksplit, std::forward<const std::vector<int64_t>>(int64s));
0045       }
0046       // If Constant node isn't used anywhere else, remove it
0047       node->removeInput(1);
0048       if (const_val->uses().size() < 1) {
0049         node_ptr->destroy();
0050       }
0051     } else {
0052       // Get Value name, find Initializer with same name
0053       for (const auto& initializer : graph->initializers()) {
0054         if (initializer.name() == inputs[1]->uniqueName()) {
0055           node->is_(ksplit, std::forward<const std::vector<int64_t>>(initializer.int64s()));
0056           node->removeInput(1);
0057           // Remove initializer
0058           if (const_val->uses().size() < 1)
0059             graph->eraseInitializerAndInput(const_val);
0060           break;
0061         }
0062       }
0063     }
0064     ONNX_ASSERTM(node->hasAttribute(ksplit), "No initializer or constant input to node found");
0065     return node;
0066   }
0067 };
0068
0069 } // namespace version_conversion
0070 } // namespace ONNX_NAMESPACE