[tf] Switch off of pywrap entry point (#14303)

This is expected to be available in 2.14 and pre-fetching so that we can
move off of more internal API end points.

Author: jpienaar

integrations/tensorflow/python_projects/iree_tf/iree/tools/tf/scripts/iree_import_tf/__main__.py

@@ -64,25 +64,30 @@ def import_saved_model(
     *, output_path, saved_model_dir, exported_names, import_type, tags
 ):
     # From here there be dragons.
-    from tensorflow.python import pywrap_mlir
+    from tensorflow.mlir.experimental import (
+        convert_saved_model_to_mlir,
+        convert_saved_model_v1_to_mlir,
+        run_pass_pipeline,
+        write_bytecode,
+    )
 
     if import_type == "savedmodel_v2":
-        result = pywrap_mlir.experimental_convert_saved_model_to_mlir(
+        result = convert_saved_model_to_mlir(
             saved_model_dir, exported_names=exported_names, show_debug_info=False
         )
     elif import_type == "savedmodel_v1":
         # You saw it here, folks: The TF team just adds random positional params
         # without explanation or default. So we detect and default them on our
         # own. Because this is normal and fine.
-        sig = inspect.signature(pywrap_mlir.experimental_convert_saved_model_v1_to_mlir)
+        sig = inspect.signature(convert_saved_model_v1_to_mlir)
         dumb_extra_kwargs = {}
         if "include_variables_in_initializers" in sig.parameters:
             dumb_extra_kwargs["include_variables_in_initializers"] = False
         if "upgrade_legacy" in sig.parameters:
             dumb_extra_kwargs["upgrade_legacy"] = False
         if "lift_variables" in sig.parameters:
             dumb_extra_kwargs["lift_variables"] = True
-        result = pywrap_mlir.experimental_convert_saved_model_v1_to_mlir(
+        result = convert_saved_model_v1_to_mlir(
             saved_model_dir,
             exported_names=exported_names,
             tags=tags,
@@ -97,15 +102,13 @@ def import_saved_model(
     # This is fine and normal, and totally to be expected. :(
     result = re.sub(r"func @__inference_(.+)_[0-9]+\(", r"func @\1(", result)
     pipeline = ["tf-lower-to-mlprogram-and-hlo"]
-    result = pywrap_mlir.experimental_run_pass_pipeline(
-        result, ",".join(pipeline), show_debug_info=False
-    )
+    result = run_pass_pipeline(result, ",".join(pipeline), show_debug_info=False)
 
-    # TODO: The experimental_write_bytecode function does not register the
+    # TODO: The write_bytecode function does not register the
     # stablehlo dialect. Once fixed, remove this bypass.
     WRITE_BYTECODE = False
     if WRITE_BYTECODE:
-        result = pywrap_mlir.experimental_write_bytecode(output_path, result)
+        result = write_bytecode(output_path, result)
     else:
         with open(output_path, "wt") as f:
             f.write(result)

integrations/tensorflow/python_projects/iree_tflite/iree/tools/tflite/scripts/iree_import_tflite/__main__.py

@@ -11,7 +11,7 @@
 import re
 import sys
 import iree.tools.tflite
-from tensorflow.python.pywrap_mlir import experimental_tflite_to_tosa_bytecode
+from tensorflow.mlir.experimental import tflite_to_tosa_bytecode
 
 
 def main():
@@ -66,7 +66,7 @@ def tflite_to_tosa(
     ordered_input_arrays=None,
     ordered_output_arrays=None,
 ):
-    experimental_tflite_to_tosa_bytecode(
+    tflite_to_tosa_bytecode(
         flatbuffer,
         bytecode,
         use_external_constant,