1# Copyright 2019 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"""Wraps toco interface with python lazy loader.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20# We need to import pywrap_tensorflow prior to the toco wrapper. 21# pylint: disable=invalid-import-order,g-bad-import-order 22from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import 23from tensorflow.python import _pywrap_toco_api 24 25# TODO(b/137402359): Remove lazy loading wrapper 26 27 28def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str, 29 debug_info_str, enable_mlir_converter): 30 """Wraps TocoConvert with lazy loader.""" 31 return _pywrap_toco_api.TocoConvert( 32 model_flags_str, 33 toco_flags_str, 34 input_data_str, 35 False, # extended_return 36 debug_info_str, 37 enable_mlir_converter) 38 39 40def wrapped_experimental_mlir_quantize(input_data_str, disable_per_channel, 41 fully_quantize, inference_type, 42 input_data_type, output_data_type, 43 enable_numeric_verify, 44 enable_whole_model_verify, 45 denylisted_ops, denylisted_nodes): 46 """Wraps experimental mlir quantize model.""" 47 return _pywrap_toco_api.ExperimentalMlirQuantizeModel( 48 input_data_str, disable_per_channel, fully_quantize, inference_type, 49 input_data_type, output_data_type, enable_numeric_verify, 50 enable_whole_model_verify, denylisted_ops, denylisted_nodes) 51 52 53def wrapped_experimental_mlir_sparsify(input_data_str): 54 """Wraps experimental mlir sparsify model.""" 55 return _pywrap_toco_api.ExperimentalMlirSparsifyModel(input_data_str) 56 57 58def wrapped_register_custom_opdefs(custom_opdefs_list): 59 """Wraps RegisterCustomOpdefs with lazy loader.""" 60 return _pywrap_toco_api.RegisterCustomOpdefs(custom_opdefs_list) 61 62 63def wrapped_retrieve_collected_errors(): 64 """Wraps RetrieveCollectedErrors with lazy loader.""" 65 return _pywrap_toco_api.RetrieveCollectedErrors() 66