From f1a2ce4406c6bb40ce0ca06440c8158313eeca6a Mon Sep 17 00:00:00 2001
From: Andrew Selle <aselle@google.com>
Date: Fri, 12 Jul 2019 17:02:17 -0700
Subject: [PATCH] Move TOCO python bindings to be built under the main
 TensorFlow Python DSO.

This allows future linking against MLIR without bloating binary. This
was not done initially since TF lite was in contrib.

PiperOrigin-RevId: 257895424
---
 tensorflow/lite/python/BUILD                  |   2 +-
 .../lite/python/interpreter_wrapper/BUILD     |   5 +-
 .../lite/python/interpreter_wrapper/numpy.cc  | 149 ++++++++++++++++++
 .../lite/python/interpreter_wrapper/numpy.h   |  15 ++
 .../interpreter_wrapper/python_utils.cc       | 146 -----------------
 .../python/interpreter_wrapper/python_utils.h |  12 +-
 tensorflow/lite/python/wrap_toco.py           |  15 +-
 tensorflow/lite/toco/python/BUILD             |  24 +--
 .../lite/toco/python/tensorflow_wrap_toco.py  |  24 +++
 tensorflow/lite/toco/python/toco.i            |   6 +
 .../lite/toco/python/toco_from_protos.py      |   4 +-
 tensorflow/python/BUILD                       |   2 +
 tensorflow/python/tensorflow.i                |   2 +
 13 files changed, 224 insertions(+), 182 deletions(-)
 create mode 100644 tensorflow/lite/toco/python/tensorflow_wrap_toco.py

diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD
index 02ccf5cc751..4b379c90980 100644
--- a/tensorflow/lite/python/BUILD
+++ b/tensorflow/lite/python/BUILD
@@ -174,7 +174,7 @@ py_library(
         "wrap_toco.py",
     ],
     deps = [
-        "//tensorflow/lite/toco/python:tensorflow_wrap_toco",
+        "//tensorflow/python:pywrap_tensorflow",
         "//tensorflow/python:util",
     ],
 )
diff --git a/tensorflow/lite/python/interpreter_wrapper/BUILD b/tensorflow/lite/python/interpreter_wrapper/BUILD
index 7209c081cb3..476f9390e57 100644
--- a/tensorflow/lite/python/interpreter_wrapper/BUILD
+++ b/tensorflow/lite/python/interpreter_wrapper/BUILD
@@ -10,6 +10,8 @@ cc_library(
     srcs = ["numpy.cc"],
     hdrs = ["numpy.h"],
     deps = [
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/c:c_api_internal",
         "//third_party/py/numpy:headers",
         "//third_party/python_runtime:headers",
     ],
@@ -27,7 +29,6 @@ cc_library(
         "//tensorflow/lite:string_util",
         "//tensorflow/lite/c:c_api_internal",
         "//tensorflow/lite/kernels:builtin_ops",
-        "//third_party/py/numpy:headers",
         "//third_party/python_runtime:headers",
         "@com_google_absl//absl/memory",
     ],
@@ -48,10 +49,8 @@ cc_library(
     srcs = ["python_utils.cc"],
     hdrs = ["python_utils.h"],
     deps = [
-        ":numpy",
         "//tensorflow/lite:framework",
         "//tensorflow/lite:string_util",
-        "//third_party/py/numpy:headers",
         "//third_party/python_runtime:headers",
     ],
 )
diff --git a/tensorflow/lite/python/interpreter_wrapper/numpy.cc b/tensorflow/lite/python/interpreter_wrapper/numpy.cc
index ff5403d2a60..d41de39210a 100644
--- a/tensorflow/lite/python/interpreter_wrapper/numpy.cc
+++ b/tensorflow/lite/python/interpreter_wrapper/numpy.cc
@@ -16,10 +16,159 @@ limitations under the License.
 #define TFLITE_IMPORT_NUMPY  // See numpy.h for explanation.
 #include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
 
+#include <memory>
+
 namespace tflite {
 namespace python {
 
 void ImportNumpy() { import_array1(); }
 
 }  // namespace python
+
+namespace python_utils {
+
+struct PyObjectDereferencer {
+  void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
+};
+using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
+
+int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
+  switch (tf_lite_type) {
+    case kTfLiteFloat32:
+      return NPY_FLOAT32;
+    case kTfLiteFloat16:
+      return NPY_FLOAT16;
+    case kTfLiteInt32:
+      return NPY_INT32;
+    case kTfLiteInt16:
+      return NPY_INT16;
+    case kTfLiteUInt8:
+      return NPY_UINT8;
+    case kTfLiteInt8:
+      return NPY_INT8;
+    case kTfLiteInt64:
+      return NPY_INT64;
+    case kTfLiteString:
+      return NPY_STRING;
+    case kTfLiteBool:
+      return NPY_BOOL;
+    case kTfLiteComplex64:
+      return NPY_COMPLEX64;
+    case kTfLiteNoType:
+      return NPY_NOTYPE;
+      // Avoid default so compiler errors created when new types are made.
+  }
+  return NPY_NOTYPE;
+}
+
+TfLiteType TfLiteTypeFromPyType(int py_type) {
+  switch (py_type) {
+    case NPY_FLOAT32:
+      return kTfLiteFloat32;
+    case NPY_INT32:
+      return kTfLiteInt32;
+    case NPY_INT16:
+      return kTfLiteInt16;
+    case NPY_UINT8:
+      return kTfLiteUInt8;
+    case NPY_INT8:
+      return kTfLiteInt8;
+    case NPY_INT64:
+      return kTfLiteInt64;
+    case NPY_BOOL:
+      return kTfLiteBool;
+    case NPY_OBJECT:
+    case NPY_STRING:
+    case NPY_UNICODE:
+      return kTfLiteString;
+    case NPY_COMPLEX64:
+      return kTfLiteComplex64;
+      // Avoid default so compiler errors created when new types are made.
+  }
+  return kTfLiteNoType;
+}
+
+TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
+  int pyarray_type = PyArray_TYPE(array);
+  return TfLiteTypeFromPyType(pyarray_type);
+}
+
+#if PY_VERSION_HEX >= 0x03030000
+bool FillStringBufferFromPyUnicode(PyObject* value,
+                                   DynamicBuffer* dynamic_buffer) {
+  Py_ssize_t len = -1;
+  const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
+  if (buf == NULL) {
+    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+#else
+bool FillStringBufferFromPyUnicode(PyObject* value,
+                                   DynamicBuffer* dynamic_buffer) {
+  UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
+  if (!utemp) {
+    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
+    return false;
+  }
+  char* buf = nullptr;
+  Py_ssize_t len = -1;
+  if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
+    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+#endif
+
+bool FillStringBufferFromPyString(PyObject* value,
+                                  DynamicBuffer* dynamic_buffer) {
+  if (PyUnicode_Check(value)) {
+    return FillStringBufferFromPyUnicode(value, dynamic_buffer);
+  }
+
+  char* buf = nullptr;
+  Py_ssize_t len = -1;
+  if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
+    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+
+bool FillStringBufferWithPyArray(PyObject* value,
+                                 DynamicBuffer* dynamic_buffer) {
+  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
+  switch (PyArray_TYPE(array)) {
+    case NPY_OBJECT:
+    case NPY_STRING:
+    case NPY_UNICODE: {
+      UniquePyObjectRef iter(PyArray_IterNew(value));
+      while (PyArray_ITER_NOTDONE(iter.get())) {
+        UniquePyObjectRef item(PyArray_GETITEM(
+            array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
+
+        if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
+          return false;
+        }
+
+        PyArray_ITER_NEXT(iter.get());
+      }
+      return true;
+    }
+    default:
+      break;
+  }
+
+  PyErr_Format(PyExc_ValueError,
+               "Cannot use numpy array of type %d for string tensor.",
+               PyArray_TYPE(array));
+  return false;
+}
+
+}  // namespace python_utils
 }  // namespace tflite
diff --git a/tensorflow/lite/python/interpreter_wrapper/numpy.h b/tensorflow/lite/python/interpreter_wrapper/numpy.h
index a3b013fcb27..98caa64a631 100644
--- a/tensorflow/lite/python/interpreter_wrapper/numpy.h
+++ b/tensorflow/lite/python/interpreter_wrapper/numpy.h
@@ -50,6 +50,8 @@ limitations under the License.
 
 #include "numpy/arrayobject.h"
 #include "numpy/ufuncobject.h"
+#include "tensorflow/lite/c/c_api_internal.h"
+#include "tensorflow/lite/string_util.h"
 
 namespace tflite {
 namespace python {
@@ -57,6 +59,19 @@ namespace python {
 void ImportNumpy();
 
 }  // namespace python
+
+namespace python_utils {
+
+int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
+
+TfLiteType TfLiteTypeFromPyType(int py_type);
+
+TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
+
+bool FillStringBufferWithPyArray(PyObject* value,
+                                 DynamicBuffer* dynamic_buffer);
+
+}  // namespace python_utils
 }  // namespace tflite
 
 #endif  // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_
diff --git a/tensorflow/lite/python/interpreter_wrapper/python_utils.cc b/tensorflow/lite/python/interpreter_wrapper/python_utils.cc
index 110c3ac4e04..1b31a0dcb54 100644
--- a/tensorflow/lite/python/interpreter_wrapper/python_utils.cc
+++ b/tensorflow/lite/python/interpreter_wrapper/python_utils.cc
@@ -17,155 +17,9 @@ limitations under the License.
 
 #include <memory>
 
-#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
-
 namespace tflite {
 namespace python_utils {
 
-struct PyObjectDereferencer {
-  void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
-};
-
-using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
-
-int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
-  switch (tf_lite_type) {
-    case kTfLiteFloat32:
-      return NPY_FLOAT32;
-    case kTfLiteFloat16:
-      return NPY_FLOAT16;
-    case kTfLiteInt32:
-      return NPY_INT32;
-    case kTfLiteInt16:
-      return NPY_INT16;
-    case kTfLiteUInt8:
-      return NPY_UINT8;
-    case kTfLiteInt8:
-      return NPY_INT8;
-    case kTfLiteInt64:
-      return NPY_INT64;
-    case kTfLiteString:
-      return NPY_STRING;
-    case kTfLiteBool:
-      return NPY_BOOL;
-    case kTfLiteComplex64:
-      return NPY_COMPLEX64;
-    case kTfLiteNoType:
-      return NPY_NOTYPE;
-      // Avoid default so compiler errors created when new types are made.
-  }
-  return NPY_NOTYPE;
-}
-
-TfLiteType TfLiteTypeFromPyType(int py_type) {
-  switch (py_type) {
-    case NPY_FLOAT32:
-      return kTfLiteFloat32;
-    case NPY_INT32:
-      return kTfLiteInt32;
-    case NPY_INT16:
-      return kTfLiteInt16;
-    case NPY_UINT8:
-      return kTfLiteUInt8;
-    case NPY_INT8:
-      return kTfLiteInt8;
-    case NPY_INT64:
-      return kTfLiteInt64;
-    case NPY_BOOL:
-      return kTfLiteBool;
-    case NPY_OBJECT:
-    case NPY_STRING:
-    case NPY_UNICODE:
-      return kTfLiteString;
-    case NPY_COMPLEX64:
-      return kTfLiteComplex64;
-      // Avoid default so compiler errors created when new types are made.
-  }
-  return kTfLiteNoType;
-}
-
-TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
-  int pyarray_type = PyArray_TYPE(array);
-  return TfLiteTypeFromPyType(pyarray_type);
-}
-
-#if PY_VERSION_HEX >= 0x03030000
-bool FillStringBufferFromPyUnicode(PyObject* value,
-                                   DynamicBuffer* dynamic_buffer) {
-  Py_ssize_t len = -1;
-  const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
-  if (buf == NULL) {
-    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
-    return false;
-  }
-  dynamic_buffer->AddString(buf, len);
-  return true;
-}
-#else
-bool FillStringBufferFromPyUnicode(PyObject* value,
-                                   DynamicBuffer* dynamic_buffer) {
-  UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
-  if (!utemp) {
-    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
-    return false;
-  }
-  char* buf = nullptr;
-  Py_ssize_t len = -1;
-  if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
-    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
-    return false;
-  }
-  dynamic_buffer->AddString(buf, len);
-  return true;
-}
-#endif
-
-bool FillStringBufferFromPyString(PyObject* value,
-                                  DynamicBuffer* dynamic_buffer) {
-  if (PyUnicode_Check(value)) {
-    return FillStringBufferFromPyUnicode(value, dynamic_buffer);
-  }
-
-  char* buf = nullptr;
-  Py_ssize_t len = -1;
-  if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
-    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
-    return false;
-  }
-  dynamic_buffer->AddString(buf, len);
-  return true;
-}
-
-bool FillStringBufferWithPyArray(PyObject* value,
-                                 DynamicBuffer* dynamic_buffer) {
-  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
-  switch (PyArray_TYPE(array)) {
-    case NPY_OBJECT:
-    case NPY_STRING:
-    case NPY_UNICODE: {
-      UniquePyObjectRef iter(PyArray_IterNew(value));
-      while (PyArray_ITER_NOTDONE(iter.get())) {
-        UniquePyObjectRef item(PyArray_GETITEM(
-            array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
-
-        if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
-          return false;
-        }
-
-        PyArray_ITER_NEXT(iter.get());
-      }
-      return true;
-    }
-    default:
-      break;
-  }
-
-  PyErr_Format(PyExc_ValueError,
-               "Cannot use numpy array of type %d for string tensor.",
-               PyArray_TYPE(array));
-  return false;
-}
-
 int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length) {
 #if PY_MAJOR_VERSION >= 3
   return PyBytes_AsStringAndSize(obj, data, length);
diff --git a/tensorflow/lite/python/interpreter_wrapper/python_utils.h b/tensorflow/lite/python/interpreter_wrapper/python_utils.h
index 7f46f0f1dad..d16caf06bf7 100644
--- a/tensorflow/lite/python/interpreter_wrapper/python_utils.h
+++ b/tensorflow/lite/python/interpreter_wrapper/python_utils.h
@@ -16,8 +16,9 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
 #define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
 
+#include <Python.h>
+
 #include "tensorflow/lite/context.h"
-#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
@@ -27,15 +28,6 @@ struct PyDecrefDeleter {
   void operator()(PyObject* p) const { Py_DECREF(p); }
 };
 
-int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
-
-TfLiteType TfLiteTypeFromPyType(int py_type);
-
-TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
-
-bool FillStringBufferWithPyArray(PyObject* value,
-                                 DynamicBuffer* dynamic_buffer);
-
 int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length);
 PyObject* ConvertToPyString(const char* data, size_t length);
 
diff --git a/tensorflow/lite/python/wrap_toco.py b/tensorflow/lite/python/wrap_toco.py
index aa17e2ff192..e9b0176f981 100644
--- a/tensorflow/lite/python/wrap_toco.py
+++ b/tensorflow/lite/python/wrap_toco.py
@@ -17,22 +17,15 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.util.lazy_loader import LazyLoader
+from tensorflow.python import pywrap_tensorflow
 
-
-# TODO(b/131123224): Lazy load since some of the performance benchmark skylark
-# rules and monolithic build break dependencies.
-_toco_python = LazyLoader(
-    "tensorflow_wrap_toco", globals(),
-    "tensorflow.lite.toco.python."
-    "tensorflow_wrap_toco")
-del LazyLoader
+# TODO(b/137402359): Remove lazy loading wrapper
 
 
 def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str,
                          debug_info_str, enable_mlir_converter):
   """Wraps TocoConvert with lazy loader."""
-  return _toco_python.TocoConvert(
+  return pywrap_tensorflow.TocoConvert(
       model_flags_str,
       toco_flags_str,
       input_data_str,
@@ -43,4 +36,4 @@ def wrapped_toco_convert(model_flags_str, toco_flags_str, input_data_str,
 
 def wrapped_get_potentially_supported_ops():
   """Wraps TocoGetPotentiallySupportedOps with lazy loader."""
-  return _toco_python.TocoGetPotentiallySupportedOps()
+  return pywrap_tensorflow.TocoGetPotentiallySupportedOps()
diff --git a/tensorflow/lite/toco/python/BUILD b/tensorflow/lite/toco/python/BUILD
index e11e9cf1578..ac78a1ece23 100644
--- a/tensorflow/lite/toco/python/BUILD
+++ b/tensorflow/lite/toco/python/BUILD
@@ -22,7 +22,9 @@ cc_library(
     name = "toco_python_api",
     srcs = ["toco_python_api.cc"],
     hdrs = ["toco_python_api.h"],
-    features = ["no_layering_check"],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
     deps = [
         "//third_party/python_runtime:headers",
         "//tensorflow/core:lib",
@@ -43,20 +45,17 @@ cc_library(
     }),
 )
 
-tf_py_wrap_cc(
+# Compatibility stub. Remove when internal customers moved.
+py_library(
     name = "tensorflow_wrap_toco",
-    srcs = ["toco.i"],
+    srcs = ["tensorflow_wrap_toco.py"],
     visibility = [
         "//learning/expander/pod/deep_pod/utils:__subpackages__",
         "//research/handwriting/converters/tflite:__subpackages__",
         "//tensorflow/lite:__subpackages__",
     ],
     deps = [
-        ":toco_python_api",
-        "//tensorflow/lite/toco:model_flags_proto_cc",
-        "//tensorflow/lite/toco:toco_flags_proto_cc",
-        "//third_party/python_runtime:headers",
-        "@com_google_absl//absl/strings",
+        "//tensorflow/python:pywrap_tensorflow",
     ],
 )
 
@@ -66,8 +65,8 @@ py_binary(
     python_version = "PY2",
     srcs_version = "PY2AND3",
     deps = [
-        ":tensorflow_wrap_toco",
         "//tensorflow/python:platform",
+        "//tensorflow/python:pywrap_tensorflow",
     ],
 )
 
@@ -84,3 +83,10 @@ tf_py_test(
         "no_pip",
     ],
 )
+
+exports_files(
+    ["toco.i"],
+    visibility = [
+        "//tensorflow/python:__subpackages__",
+    ],
+)
diff --git a/tensorflow/lite/toco/python/tensorflow_wrap_toco.py b/tensorflow/lite/toco/python/tensorflow_wrap_toco.py
new file mode 100644
index 00000000000..d70b0438886
--- /dev/null
+++ b/tensorflow/lite/toco/python/tensorflow_wrap_toco.py
@@ -0,0 +1,24 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Stub to make toco convert accessible."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# TODO(aselle): Remove once no clients internally depend on this.
+
+# pylint: disable=unused-import
+from tensorflow.python.pywrap_tensorflow import TocoConvert
+# pylint: enable=unused-import
diff --git a/tensorflow/lite/toco/python/toco.i b/tensorflow/lite/toco/python/toco.i
index 3aa9ce6553f..1016091ed77 100644
--- a/tensorflow/lite/toco/python/toco.i
+++ b/tensorflow/lite/toco/python/toco.i
@@ -19,6 +19,12 @@ limitations under the License.
 #include "tensorflow/lite/toco/python/toco_python_api.h"
 %}
 
+// The TensorFlow exception handler releases the GIL with
+// Py_BEGIN_ALLOW_THREADS. Remove that because these function use the Python
+// API to decode inputs.
+%noexception toco::TocoConvert;
+%noexception toco::TocoGetPotentiallySupportedOps;
+
 namespace toco {
 
 // Convert a model represented in `input_contents`. `model_flags_proto`
diff --git a/tensorflow/lite/toco/python/toco_from_protos.py b/tensorflow/lite/toco/python/toco_from_protos.py
index 0566cb8ba60..5137ea165b3 100644
--- a/tensorflow/lite/toco/python/toco_from_protos.py
+++ b/tensorflow/lite/toco/python/toco_from_protos.py
@@ -19,7 +19,7 @@ from __future__ import print_function
 
 import argparse
 import sys
-from tensorflow.lite.toco.python import tensorflow_wrap_toco
+from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.platform import app
 
 FLAGS = None
@@ -43,7 +43,7 @@ def execute(unused_args):
 
   enable_mlir_converter = FLAGS.enable_mlir_converter
 
-  output_str = tensorflow_wrap_toco.TocoConvert(
+  output_str = pywrap_tensorflow.TocoConvert(
       model_str,
       toco_str,
       input_str,
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 5e5836f1d82..c1348c7f992 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -4759,6 +4759,7 @@ tf_py_wrap_cc(
         "util/tfprof.i",
         "util/transform_graph.i",
         "util/util.i",
+        "//tensorflow/lite/toco/python:toco.i",
     ],
     # add win_def_file for pywrap_tensorflow
     win_def_file = select({
@@ -4805,6 +4806,7 @@ tf_py_wrap_cc(
         "//tensorflow/core/distributed_runtime:server_lib",
         "//tensorflow/core/profiler/internal:print_model_analysis",
         "//tensorflow/tools/graph_transforms:transform_graph_lib",
+        "//tensorflow/lite/toco/python:toco_python_api",
         "//tensorflow/python/eager:pywrap_tfe_lib",
     ] + (tf_additional_lib_deps() +
          tf_additional_plugin_deps() +
diff --git a/tensorflow/python/tensorflow.i b/tensorflow/python/tensorflow.i
index 4e1bf3d8362..82e30be33a3 100644
--- a/tensorflow/python/tensorflow.i
+++ b/tensorflow/python/tensorflow.i
@@ -36,6 +36,8 @@ limitations under the License.
 
 %include "tensorflow/python/lib/core/bfloat16.i"
 
+%include "tensorflow/lite/toco/python/toco.i"
+
 %include "tensorflow/python/lib/io/file_io.i"
 %include "tensorflow/python/training/quantize_training.i"