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
26# TODO(b/137402359): Remove lazy loading wrapper
27
28
29def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str,
30                         debug_info_str, enable_mlir_converter):
31  """Wraps TocoConvert with lazy loader."""
32  return _pywrap_toco_api.TocoConvert(
33      model_flags_str,
34      toco_flags_str,
35      input_data_str,
36      False,  # extended_return
37      debug_info_str,
38      enable_mlir_converter)
39
40
41def wrapped_get_potentially_supported_ops():
42  """Wraps TocoGetPotentiallySupportedOps with lazy loader."""
43  return _pywrap_toco_api.TocoGetPotentiallySupportedOps()
44
45
46def wrapped_experimental_mlir_quantize(input_data_str, disable_per_channel,
47                                       fully_quantize, inference_type,
48                                       enable_numeric_verify):
49  """Wraps experimental mlir quantize model."""
50  return _pywrap_toco_api.ExperimentalMlirQuantizeModel(input_data_str,
51                                                        disable_per_channel,
52                                                        fully_quantize,
53                                                        inference_type,
54                                                        enable_numeric_verify)
55
56
57def wrapped_experimental_mlir_sparsify(input_data_str):
58  """Wraps experimental mlir sparsify model."""
59  return _pywrap_toco_api.ExperimentalMlirSparsifyModel(input_data_str)
60
61
62def wrapped_register_custom_opdefs(custom_opdefs_list):
63  """Wraps RegisterCustomOpdefs with lazy loader."""
64  return _pywrap_toco_api.RegisterCustomOpdefs(custom_opdefs_list)
65