1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #include <memory> 16 17 #include "tensorflow/lite/kernels/register.h" 18 19 #include "tensorflow_lite_support/custom_ops/kernel/unsorted_segment.h" 20 21 namespace tflite { 22 namespace task { 23 // Create a custom op resolver to provide the unsorted_segment_prod op 24 // required by the bert_nl_classifier and rb_model for BertNLClassifier. CreateOpResolver()25std::unique_ptr<tflite::OpResolver> CreateOpResolver() { // NOLINT 26 std::unique_ptr<tflite::ops::builtin::BuiltinOpResolver> resolver( 27 new tflite::ops::builtin::BuiltinOpResolver); 28 // "UnsortedSegmentProd" is the name used by unsorted_segment_prod op when 29 // when converting SavedModel to tflite using the size optimization approach. 30 resolver->AddCustom("UnsortedSegmentProd", 31 tflite::ops::custom::Register_UNSORTED_SEGMENT_PROD()); 32 // "FlexUnsortedSegmentProd" is the name used by unsorted_segment_prod op when 33 // when converting SavedModel to tflite using the the other approaches. 34 resolver->AddCustom("FlexUnsortedSegmentProd", 35 tflite::ops::custom::Register_UNSORTED_SEGMENT_PROD()); 36 return std::unique_ptr<tflite::OpResolver>(std::move(resolver)); 37 } 38 39 } // namespace task 40 } // namespace tflite