# Copyright 2020 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Python module for MLIR functions exported by pybind11."""

# pylint: disable=invalid-import-order, g-bad-import-order, wildcard-import, unused-import, undefined-variable

from tensorflow.python import pywrap_tensorflow

from tensorflow.python.eager import context

from tensorflow.python._pywrap_mlir import *

def import_graphdef(graphdef,

pass_pipeline,

show_debug_info,

input_names=None,

input_data_types=None,

input_data_shapes=None,

output_names=[]):

if input_names is not None:

return ImportGraphDef(

str(graphdef).encode('utf-8'), pass_pipeline.encode('utf-8'),

show_debug_info, ','.join(input_names).encode('utf-8'),

','.join(input_data_types).encode('utf-8'),

':'.join(input_data_shapes).encode('utf-8'),

','.join(output_names).encode('utf-8'))

return ImportGraphDef(

str(graphdef).encode('utf-8'), pass_pipeline.encode('utf-8'),

show_debug_info)

def import_function(concrete_function, pass_pipeline, show_debug_info):

ctxt = context.context()

ctxt.ensure_initialized()

return ImportFunction(ctxt._handle,

str(concrete_function.function_def).encode('utf-8'),

pass_pipeline.encode('utf-8'), show_debug_info)

def experimental_convert_saved_model_to_mlir(saved_model_path, exported_names,

show_debug_info):

return ExperimentalConvertSavedModelToMlir(

str(saved_model_path).encode('utf-8'),

str(exported_names).encode('utf-8'), show_debug_info)

def experimental_convert_saved_model_v1_to_mlir_lite(saved_model_path,

exported_names, tags,

upgrade_legacy,

show_debug_info):

return ExperimentalConvertSavedModelV1ToMlirLite(

str(saved_model_path).encode('utf-8'),

str(exported_names).encode('utf-8'),

str(tags).encode('utf-8'), upgrade_legacy, show_debug_info)

def experimental_convert_saved_model_v1_to_mlir(saved_model_path,

exported_names, tags,

lift_variables, upgrade_legacy,

show_debug_info):

return ExperimentalConvertSavedModelV1ToMlir(

str(saved_model_path).encode('utf-8'),

str(exported_names).encode('utf-8'),

str(tags).encode('utf-8'), lift_variables, upgrade_legacy,

show_debug_info)

def experimental_run_pass_pipeline(mlir_txt, pass_pipeline, show_debug_info):

return ExperimentalRunPassPipeline(

mlir_txt.encode('utf-8'), pass_pipeline.encode('utf-8'), show_debug_info)