Rename 'pyiree' to 'iree' (#5145)
Renames all non-BUILD references to `pyiree` (`BUILD` files are changed in
an internal patch). Simplifies the imports in some places to be more consistent.
Also deletes some old test files that use the old API and enables the tests for
`iree.tf.support`.
diff --git a/integrations/tensorflow/bindings/python/CMakeLists.txt b/integrations/tensorflow/bindings/python/CMakeLists.txt
index f88027a..ee1fab2 100644
--- a/integrations/tensorflow/bindings/python/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/CMakeLists.txt
@@ -20,17 +20,17 @@
endfunction()
if(${IREE_BUILD_TENSORFLOW_COMPILER})
- _add_overlay_subdirectory(pyiree/tools/tf)
+ _add_overlay_subdirectory(iree/tools/tf)
endif()
if(${IREE_BUILD_TFLITE_COMPILER})
- _add_overlay_subdirectory(pyiree/tools/tflite)
+ _add_overlay_subdirectory(iree/tools/tflite)
endif()
if(${IREE_BUILD_XLA_COMPILER})
- _add_overlay_subdirectory(pyiree/tools/xla)
+ _add_overlay_subdirectory(iree/tools/xla)
endif()
# TODO: Find another place for the TF support library.
# Pure python so can just always be generated.
-_add_overlay_subdirectory(pyiree/tf/support)
+_add_overlay_subdirectory(iree/tf/support)
diff --git a/integrations/tensorflow/bindings/python/iree/tf/support/CMakeLists.txt b/integrations/tensorflow/bindings/python/iree/tf/support/CMakeLists.txt
new file mode 100644
index 0000000..672d898
--- /dev/null
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/CMakeLists.txt
@@ -0,0 +1,65 @@
+# Copyright 2020 Google LLC
+#
+# 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
+#
+# https://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.
+
+iree_py_library(
+ NAME
+ support
+ SRCS
+ "__init__.py"
+ "module_utils.py"
+ "tf_test_driver.py"
+ "tf_test_utils.py"
+ "tf_utils.py"
+ "trace_utils.py"
+)
+
+# TODO(#5147): Replace the hack below with something more thought out.
+# Make the base dir that run_test.sh expects to exist. This dir doesn't exist
+# by default because we symlink these python files into bindings/python/iree.
+file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/integrations/tensorflow/bindings/python/iree/tf/support")
+
+iree_py_test(
+ NAME
+ module_utils_test
+ SRCS
+ "module_utils_test.py"
+)
+
+iree_py_test(
+ NAME
+ tf_test_utils_test
+ SRCS
+ "tf_test_utils_test.py"
+)
+
+iree_py_test(
+ NAME
+ tf_utils_test
+ SRCS
+ "tf_utils_test.py"
+)
+
+iree_py_test(
+ NAME
+ trace_utils_test
+ SRCS
+ "trace_utils_test.py"
+)
+
+iree_py_install_package(
+ AUGMENT_EXISTING_PACKAGE
+ COMPONENT IreePythonPackage-tools-tf
+ PACKAGE_NAME iree_tools_tf
+ MODULE_PATH iree/tf/support
+)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/__init__.py b/integrations/tensorflow/bindings/python/iree/tf/support/__init__.py
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/__init__.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/__init__.py
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils.py b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
index 7f3051c..b358095 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
@@ -23,10 +23,10 @@
from absl import flags
from absl import logging
+import iree.compiler.tf
+import iree.rt
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree import rt
-from pyiree.compiler import tf as tf_compiler
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
flags.DEFINE_bool(
@@ -117,7 +117,7 @@
backend_info.backend_id,
needs_temp_saved_model_dir=True,
) if artifacts_dir else {})
- immediate_result = tf_compiler.compile_module(
+ immediate_result = iree.compiler.tf.compile_module(
module,
target_backends=backend_info.compiler_targets,
exported_names=exported_names,
@@ -154,7 +154,7 @@
"""
output_kwargs = (_get_tf_import_output_kwargs(
artifacts_dir, backend_info.backend_id) if artifacts_dir else {})
- immediate_result = tf_compiler.compile_saved_model(
+ immediate_result = iree.compiler.tf.compile_saved_model(
saved_model_dir,
import_type="SIGNATURE_DEF",
target_backends=backend_info.compiler_targets,
@@ -283,7 +283,8 @@
class _IreeFunctionWrapper(_FunctionWrapper):
"""Wraps an IREE function, making it callable."""
- def __init__(self, context: rt.SystemContext, f: rt.system_api.BoundFunction):
+ def __init__(self, context: iree.rt.SystemContext,
+ f: iree.rt.system_api.BoundFunction):
self._context = context
self._f = f
@@ -303,8 +304,8 @@
module_name: str,
backend_info: "BackendInfo",
compiled_paths: Dict[str, str],
- vm_module: rt.VmModule,
- config: rt.Config,
+ vm_module: iree.rt.VmModule,
+ config: iree.rt.Config,
):
"""Base constructor – Use one of the named constructors instead.
@@ -314,8 +315,8 @@
backend_info: BackendInfo with the details about compiling this module.
compiled_paths: A dictionary mapping compiled method names to file paths
corresponding to their serialized representations.
- vm_module: A rt.VmModule containing compilation info to wrap.
- config: A rt.Config containing compilation info to wrap.
+ vm_module: A iree.rt.VmModule containing compilation info to wrap.
+ config: A iree.rt.Config containing compilation info to wrap.
"""
super().__init__(module_name, backend_info, compiled_paths)
self._vm_module = vm_module
@@ -364,8 +365,8 @@
backend_info=backend_info,
exported_names=exported_names,
artifacts_dir=artifacts_dir)
- vm_module = rt.VmModule.from_flatbuffer(module_blob)
- config = rt.Config(driver_name=backend_info.driver)
+ vm_module = iree.rt.VmModule.from_flatbuffer(module_blob)
+ config = iree.rt.Config(driver_name=backend_info.driver)
compiled_paths = None
if compiled_path is not None:
@@ -406,8 +407,8 @@
module_blob, compiled_path = _incrementally_compile_tf_signature_def_saved_model(
saved_model_dir, saved_model_tags, backend_info, exported_name,
artifacts_dir)
- vm_module = rt.VmModule.from_flatbuffer(module_blob)
- config = rt.Config(driver_name=backend_info.driver)
+ vm_module = iree.rt.VmModule.from_flatbuffer(module_blob)
+ config = iree.rt.Config(driver_name=backend_info.driver)
compiled_paths = None
if compiled_path is not None:
@@ -420,8 +421,8 @@
"""Reinitializes all stateful variables."""
# set_random_seed is not needed here because the model_class.__init__ is not
# called.
- self._context = rt.SystemContext(modules=[self._vm_module],
- config=self._config)
+ self._context = iree.rt.SystemContext(modules=[self._vm_module],
+ config=self._config)
def __getattr__(self, attr: str) -> _IreeFunctionWrapper:
# Try to resolve it as a function.
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils_test.py b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils_test.py
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils_test.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/module_utils_test.py
index 8baa6bc..b43d39d 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils_test.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils_test.py
@@ -12,14 +12,14 @@
# 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.
-"""Tests for pyiree.tf.support.module_utils."""
+"""Tests for iree.tf.support.module_utils."""
import os
import tempfile
from absl import logging
from absl.testing import parameterized
-from pyiree.tf.support import module_utils
+from iree.tf.support import module_utils
import tensorflow as tf
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py b/integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils.py
similarity index 99%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils.py
index 24e83aa..f1ff5b0 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils.py
@@ -34,9 +34,9 @@
from absl import flags
from absl import logging
-from pyiree.tf.support import module_utils
-from pyiree.tf.support import tf_utils
-from pyiree.tf.support import trace_utils
+from iree.tf.support import module_utils
+from iree.tf.support import tf_utils
+from iree.tf.support import trace_utils
import tensorflow.compat.v2 as tf
flags.DEFINE_string("reference_backend", "tf",
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py b/integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils_test.py
similarity index 95%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils_test.py
index 84c1c36..697a6d5 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/tf_test_utils_test.py
@@ -12,11 +12,11 @@
# 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.
-"""Tests for pyiree.tf.support.tf_test_utils."""
+"""Tests for iree.tf.support.tf_test_utils."""
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow as tf
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py b/integrations/tensorflow/bindings/python/iree/tf/support/tf_utils.py
similarity index 99%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/tf_utils.py
index 2c9b478..d0fddd2 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/tf_utils.py
@@ -20,8 +20,8 @@
from typing import Any, Callable, Sequence, Set, Tuple, Union
from absl import logging
+import iree.rt
import numpy as np
-from pyiree import rt
import tensorflow.compat.v2 as tf
InputGeneratorType = Callable[[Sequence[int], Union[tf.DType, np.dtype]],
@@ -90,7 +90,7 @@
def convert_to_numpy(values: Any) -> Any:
"""Converts any tf.Tensor, int, float, bool or list values to numpy."""
- return apply_function(values, rt.normalize_value)
+ return apply_function(values, iree.rt.normalize_value)
def to_mlir_type(dtype: np.dtype) -> str:
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py b/integrations/tensorflow/bindings/python/iree/tf/support/tf_utils_test.py
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/tf_utils_test.py
index deef8d9..9da1319 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/tf_utils_test.py
@@ -12,11 +12,11 @@
# 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.
-"""Tests for pyiree.tf.support.tf_utils."""
+"""Tests for iree.tf.support.tf_utils."""
from absl.testing import parameterized
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_utils
import tensorflow as tf
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils.py b/integrations/tensorflow/bindings/python/iree/tf/support/trace_utils.py
similarity index 99%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/trace_utils.py
index e6c4265..f5ae8d3 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/trace_utils.py
@@ -30,9 +30,9 @@
from typing import Any, Callable, Dict, Sequence, Tuple, Union
from absl import logging
+from iree.tf.support import module_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import module_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
NUMPY_LINEWIDTH = 120
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils_test.py b/integrations/tensorflow/bindings/python/iree/tf/support/trace_utils_test.py
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils_test.py
rename to integrations/tensorflow/bindings/python/iree/tf/support/trace_utils_test.py
index 58315c8..c092269 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/trace_utils_test.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/trace_utils_test.py
@@ -12,15 +12,15 @@
# 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.
-"""Tests for pyiree.tf.support.trace_utils."""
+"""Tests for iree.tf.support.trace_utils."""
import os
import tempfile
from absl.testing import parameterized
+from iree.tf.support import module_utils
+from iree.tf.support import trace_utils
import numpy as np
-from pyiree.tf.support import module_utils
-from pyiree.tf.support import trace_utils
import tensorflow as tf
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt b/integrations/tensorflow/bindings/python/iree/tools/tf/CMakeLists.txt
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt
rename to integrations/tensorflow/bindings/python/iree/tools/tf/CMakeLists.txt
index 2a6fe5b..4d7792f 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/tf/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/iree/tools/tf/CMakeLists.txt
@@ -30,7 +30,7 @@
iree_py_install_package(
COMPONENT IreePythonPackage-tools-tf
PACKAGE_NAME iree_tools_tf
- MODULE_PATH pyiree/tools/tf
+ MODULE_PATH iree/tools/tf
DEPS
iree_tf_compiler_iree-tf-import
)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tf/__init__.py b/integrations/tensorflow/bindings/python/iree/tools/tf/__init__.py
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tf/__init__.py
rename to integrations/tensorflow/bindings/python/iree/tools/tf/__init__.py
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tf/setup.py.in b/integrations/tensorflow/bindings/python/iree/tools/tf/setup.py.in
similarity index 93%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tf/setup.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/tf/setup.py.in
index 3f75596..b83ddfd 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/tf/setup.py.in
+++ b/integrations/tensorflow/bindings/python/iree/tools/tf/setup.py.in
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-# Build platform specific wheel files for the pyiree.rt package.
+# Build platform specific wheel files for the iree.rt package.
# Built artifacts are per-platform and build out of the build tree.
from distutils.command.install import install
@@ -79,11 +79,11 @@
],
python_requires=">=3.6",
packages=find_namespace_packages(include=[
- "pyiree.tools.tf",
- "pyiree.tf.support",
+ "iree.tools.tf",
+ "iree.tf.support",
]),
package_data={
- "pyiree.tools.tf": [f"iree-tf-import{exe_suffix}",],
+ "iree.tools.tf": [f"iree-tf-import{exe_suffix}",],
},
cmdclass={
'bdist_wheel': bdist_wheel,
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tf/version.py.in b/integrations/tensorflow/bindings/python/iree/tools/tf/version.py.in
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tf/version.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/tf/version.py.in
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/CMakeLists.txt b/integrations/tensorflow/bindings/python/iree/tools/tflite/CMakeLists.txt
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tflite/CMakeLists.txt
rename to integrations/tensorflow/bindings/python/iree/tools/tflite/CMakeLists.txt
index b2ff756..2406091 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/iree/tools/tflite/CMakeLists.txt
@@ -30,7 +30,7 @@
iree_py_install_package(
COMPONENT IreePythonPackage-tools-tflite
PACKAGE_NAME iree_tools_tflite
- MODULE_PATH pyiree/tools/tflite
+ MODULE_PATH iree/tools/tflite
DEPS
iree_tf_compiler_iree-import-tflite
)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/__init__.py b/integrations/tensorflow/bindings/python/iree/tools/tflite/__init__.py
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tflite/__init__.py
rename to integrations/tensorflow/bindings/python/iree/tools/tflite/__init__.py
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/setup.py.in b/integrations/tensorflow/bindings/python/iree/tools/tflite/setup.py.in
similarity index 92%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tflite/setup.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/tflite/setup.py.in
index 65bd8ab..ca4c45f 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/setup.py.in
+++ b/integrations/tensorflow/bindings/python/iree/tools/tflite/setup.py.in
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-# Build platform specific wheel files for the pyiree.rt package.
+# Build platform specific wheel files for the iree.rt package.
# Built artifacts are per-platform and build out of the build tree.
from distutils.command.install import install
@@ -78,9 +78,9 @@
"Development Status :: 3 - Alpha",
],
python_requires=">=3.6",
- packages=find_namespace_packages(include=["pyiree.tools.tflite"]),
+ packages=find_namespace_packages(include=["iree.tools.tflite"]),
package_data={
- "pyiree.tools.tflite": [f"iree-import-tflite{exe_suffix}",],
+ "iree.tools.tflite": [f"iree-import-tflite{exe_suffix}",],
},
cmdclass={
'bdist_wheel': bdist_wheel,
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/tflite/version.py.in b/integrations/tensorflow/bindings/python/iree/tools/tflite/version.py.in
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/tflite/version.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/tflite/version.py.in
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt b/integrations/tensorflow/bindings/python/iree/tools/xla/CMakeLists.txt
similarity index 97%
rename from integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt
rename to integrations/tensorflow/bindings/python/iree/tools/xla/CMakeLists.txt
index 56c23ae..32a0fd7 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/xla/CMakeLists.txt
+++ b/integrations/tensorflow/bindings/python/iree/tools/xla/CMakeLists.txt
@@ -30,7 +30,7 @@
iree_py_install_package(
COMPONENT IreePythonPackage-tools-xla
PACKAGE_NAME iree_tools_xla
- MODULE_PATH pyiree/tools/xla
+ MODULE_PATH iree/tools/xla
DEPS
iree_tf_compiler_iree-import-xla
)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/__init__.py b/integrations/tensorflow/bindings/python/iree/tools/xla/__init__.py
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/xla/__init__.py
rename to integrations/tensorflow/bindings/python/iree/tools/xla/__init__.py
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/setup.py.in b/integrations/tensorflow/bindings/python/iree/tools/xla/setup.py.in
similarity index 92%
rename from integrations/tensorflow/bindings/python/pyiree/tools/xla/setup.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/xla/setup.py.in
index e863337..d604cd9 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tools/xla/setup.py.in
+++ b/integrations/tensorflow/bindings/python/iree/tools/xla/setup.py.in
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-# Build platform specific wheel files for the pyiree.rt package.
+# Build platform specific wheel files for the iree.rt package.
# Built artifacts are per-platform and build out of the build tree.
from distutils.command.install import install
@@ -78,9 +78,9 @@
"Development Status :: 3 - Alpha",
],
python_requires=">=3.6",
- packages=find_namespace_packages(include=["pyiree.tools.xla"]),
+ packages=find_namespace_packages(include=["iree.tools.xla"]),
package_data={
- "pyiree.tools.xla": [f"iree-import-xla{exe_suffix}",],
+ "iree.tools.xla": [f"iree-import-xla{exe_suffix}",],
},
cmdclass={
'bdist_wheel': bdist_wheel,
diff --git a/integrations/tensorflow/bindings/python/pyiree/tools/xla/version.py.in b/integrations/tensorflow/bindings/python/iree/tools/xla/version.py.in
similarity index 100%
rename from integrations/tensorflow/bindings/python/pyiree/tools/xla/version.py.in
rename to integrations/tensorflow/bindings/python/iree/tools/xla/version.py.in
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/CMakeLists.txt b/integrations/tensorflow/bindings/python/pyiree/tf/support/CMakeLists.txt
deleted file mode 100644
index dab701b..0000000
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/CMakeLists.txt
+++ /dev/null
@@ -1,34 +0,0 @@
-# Copyright 2020 Google LLC
-#
-# 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
-#
-# https://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.
-
-iree_py_library(
- NAME
- support
- SRCS
- "__init__.py"
- "module_utils.py"
- "tf_test_driver.py"
- "tf_test_utils.py"
- "tf_utils.py"
- "trace_utils.py"
-)
-
-# TODO: Add tests.
-
-iree_py_install_package(
- AUGMENT_EXISTING_PACKAGE
- COMPONENT IreePythonPackage-tools-tf
- PACKAGE_NAME iree_tools_tf
- MODULE_PATH pyiree/tf/support
-)
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_driver.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_driver.py
deleted file mode 100644
index 44e7906..0000000
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_driver.py
+++ /dev/null
@@ -1,129 +0,0 @@
-# Lint as: python3
-# Copyright 2019 Google LLC
-#
-# 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
-#
-# https://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.
-"""Utilities for running tests from TensorFlow models."""
-
-import contextlib
-import io
-import os
-import subprocess
-import sys
-import tempfile
-import traceback
-
-from absl import app
-from absl import flags
-from pyiree.tf import compiler
-import tensorflow.compat.v2 as tf
-
-flags.DEFINE_string("filecheck_binary", "filecheck",
- "Location of the filecheck binary.")
-flags.DEFINE_bool("disable_filecheck", False,
- "Disables filecheck redirection (for debugging).")
-FLAGS = flags.FLAGS
-
-ALL_TEST_DICTS = []
-
-
-def add_test(**kwargs):
- assert "test_name" in kwargs, "'test_name' is a required argument"
- ALL_TEST_DICTS.append(kwargs)
-
-
-def _run_test(test_dict):
- """Runs an individual test dict."""
- tf_module_builder_lambda = test_dict["tf_module_builder"]
- tf_module = tf_module_builder_lambda()
- input_module = compiler.tf_module_to_compiler_module(tf_module,
- pass_pipeline=())
-
- passes = test_dict.get("passes")
- expect_pass_failure = test_dict.get("expect_pass_failure")
- if passes:
- try:
- input_module.run_pass_pipeline(passes)
- except: # pylint: disable=bare-except
- if not expect_pass_failure:
- print("UNEXPECTED PASS FAILURE (INTERMEDIATE ASM FOLLOWS ON STDERR):",
- file=sys.stderr)
- print(input_module.to_asm(), file=sys.stderr)
- raise
-
- # Print the input module ASM.
- if test_dict.get("print_input_module"):
- print(input_module.to_asm())
-
-
-def _internal_run_tests():
- """Main function that runs all tests."""
- test_count = 0
- for test_dict in ALL_TEST_DICTS:
- test_count += 1
- test_name = test_dict["test_name"]
- print("RUN_TEST:", test_name)
- try:
- _run_test(test_dict)
- print("FINISH_TEST:", test_name)
- except: # pylint: disable=bare-except
- # Error goes to stdout for FileCheck.
- traceback.print_exc(file=sys.stdout)
- print("FINISH_TEST_WITH_EXCEPTION:", test_name)
-
- print("FINISHED: RAN", test_count, "TESTS", file=sys.stderr)
-
-
-def _find_filecheck():
- """Finds the filecheck binary."""
- filecheck_binary = FLAGS.filecheck_binary
- if os.path.isabs(filecheck_binary):
- return filecheck_binary
- # TODO(laurenzo): Why is this runfiles resolution so hard and undocumented.
- # Talk to bazel team.
- runfiles_dir = os.environ.get("RUNFILES_DIR")
- if runfiles_dir:
- workspace_name = os.environ.get("TEST_WORKSPACE")
- if workspace_name:
- runfiles_dir = os.path.join(runfiles_dir, workspace_name)
- filecheck_binary = os.path.join(runfiles_dir, filecheck_binary)
- # Convert forward slash version to platform default (Windows).
- filecheck_binary = filecheck_binary.replace("/", os.path.sep)
- return filecheck_binary
-
-
-def run_tests(main_file, with_filecheck=True):
- """Main entry point."""
-
- def internal_main(unused_argv):
- """App main."""
- # In case if running with a version prior to v2 defaulting.
- tf.enable_v2_behavior()
- if with_filecheck and not FLAGS.disable_filecheck:
- # Capture and run through filecheck.
- filecheck_capture_io = io.StringIO()
- with contextlib.redirect_stdout(filecheck_capture_io):
- _internal_run_tests()
- filecheck_capture_io.flush()
- filecheck_input = filecheck_capture_io.getvalue()
- # Convert forward slash version to platform default (Windows).
- filecheck_binary = _find_filecheck()
- filecheck_args = [filecheck_binary, main_file, "--dump-input=fail"]
- print("LAUNCHING FILECHECK:", filecheck_args, file=sys.stderr)
- p = subprocess.Popen(filecheck_args, stdin=subprocess.PIPE)
- p.communicate(filecheck_input.encode("UTF-8"))
- sys.exit(p.returncode)
- else:
- # Just run directly.
- _internal_run_tests()
-
- app.run(internal_main)
diff --git a/integrations/tensorflow/e2e/README.md b/integrations/tensorflow/e2e/README.md
index 1210edb..fb7ebd4 100644
--- a/integrations/tensorflow/e2e/README.md
+++ b/integrations/tensorflow/e2e/README.md
@@ -40,13 +40,13 @@
specified directory. These artifacts include MLIR across various lowerings and
the compiled VM FlatBuffer. A basic example of creating and calling an
`IreeCompiledModule` can be found in
-[`module_utils_test.py`](https://github.com/google/iree/blob/main/integrations/tensorflow/bindings/python/pyiree/tf/support/module_utils_test.py)
+[`module_utils_test.py`](https://github.com/google/iree/blob/main/integrations/tensorflow/bindings/python/iree/tf/support/module_utils_test.py)
When using Keras models or tf.Modules with functions that IREE can't compile,
`exported_names` should be specified. For example:
```python
-from pyiree.tf.support import module_utils
+from iree.tf.support import module_utils
vmla_module = module_utils.IreeCompiledModule(
module_class=KerasTFModuleClass,
backend_info=module_utils.BackendInfo('iree_vmla'),
diff --git a/integrations/tensorflow/e2e/batch_norm_test.py b/integrations/tensorflow/e2e/batch_norm_test.py
index 13b22d9..44dd8cf 100644
--- a/integrations/tensorflow/e2e/batch_norm_test.py
+++ b/integrations/tensorflow/e2e/batch_norm_test.py
@@ -15,9 +15,9 @@
"""Batch norm tests."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/batch_to_space_nd_test.py b/integrations/tensorflow/e2e/batch_to_space_nd_test.py
index 7c4120e..7671a82 100644
--- a/integrations/tensorflow/e2e/batch_to_space_nd_test.py
+++ b/integrations/tensorflow/e2e/batch_to_space_nd_test.py
@@ -15,9 +15,9 @@
"""Batch To Space ND tests."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/broadcast_to_test.py b/integrations/tensorflow/e2e/broadcast_to_test.py
index 7eb0936..9d4e170 100644
--- a/integrations/tensorflow/e2e/broadcast_to_test.py
+++ b/integrations/tensorflow/e2e/broadcast_to_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/broadcasting_test.py b/integrations/tensorflow/e2e/broadcasting_test.py
index df2d387..a502666 100644
--- a/integrations/tensorflow/e2e/broadcasting_test.py
+++ b/integrations/tensorflow/e2e/broadcasting_test.py
@@ -15,9 +15,9 @@
"""Test broadcasting support."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/concat_test.py b/integrations/tensorflow/e2e/concat_test.py
index 419ebe9..0288e64 100644
--- a/integrations/tensorflow/e2e/concat_test.py
+++ b/integrations/tensorflow/e2e/concat_test.py
@@ -15,9 +15,9 @@
"""Test concat op."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/control_flow_test.py b/integrations/tensorflow/e2e/control_flow_test.py
index 180f2da..a69867e 100644
--- a/integrations/tensorflow/e2e/control_flow_test.py
+++ b/integrations/tensorflow/e2e/control_flow_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/conv_test.py b/integrations/tensorflow/e2e/conv_test.py
index ca0076e..b2c9236 100644
--- a/integrations/tensorflow/e2e/conv_test.py
+++ b/integrations/tensorflow/e2e/conv_test.py
@@ -14,9 +14,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/conv_transpose_test.py b/integrations/tensorflow/e2e/conv_transpose_test.py
index f62eab4..cfe6471 100644
--- a/integrations/tensorflow/e2e/conv_transpose_test.py
+++ b/integrations/tensorflow/e2e/conv_transpose_test.py
@@ -14,9 +14,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/depth_conv_test.py b/integrations/tensorflow/e2e/depth_conv_test.py
index 1a73413..2905ba4 100644
--- a/integrations/tensorflow/e2e/depth_conv_test.py
+++ b/integrations/tensorflow/e2e/depth_conv_test.py
@@ -14,9 +14,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/dynamic_mlp_relu_test.py b/integrations/tensorflow/e2e/dynamic_mlp_relu_test.py
index 87b1636..7d5086a 100644
--- a/integrations/tensorflow/e2e/dynamic_mlp_relu_test.py
+++ b/integrations/tensorflow/e2e/dynamic_mlp_relu_test.py
@@ -18,9 +18,9 @@
# (unimplemented dynamic dot_general).
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
HIDDEN_1_DIM = 256
diff --git a/integrations/tensorflow/e2e/dynamic_mlp_test.py b/integrations/tensorflow/e2e/dynamic_mlp_test.py
index 4a531e6..e1ef21a 100644
--- a/integrations/tensorflow/e2e/dynamic_mlp_test.py
+++ b/integrations/tensorflow/e2e/dynamic_mlp_test.py
@@ -14,9 +14,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
HIDDEN_1_DIM = 256
diff --git a/integrations/tensorflow/e2e/einsum_dynamic_test.py b/integrations/tensorflow/e2e/einsum_dynamic_test.py
index 2904da4..795e640 100644
--- a/integrations/tensorflow/e2e/einsum_dynamic_test.py
+++ b/integrations/tensorflow/e2e/einsum_dynamic_test.py
@@ -14,8 +14,8 @@
# limitations under the License.
"""Test matrix ops via einsum"""
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
LEFT_DIM = 6
diff --git a/integrations/tensorflow/e2e/einsum_static_test.py b/integrations/tensorflow/e2e/einsum_static_test.py
index 316d8be..6709215 100644
--- a/integrations/tensorflow/e2e/einsum_static_test.py
+++ b/integrations/tensorflow/e2e/einsum_static_test.py
@@ -14,8 +14,8 @@
# limitations under the License.
"""Test matrix ops via einsum"""
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
LEFT_DIM = 6
diff --git a/integrations/tensorflow/e2e/einsum_vector_test.py b/integrations/tensorflow/e2e/einsum_vector_test.py
index 3f64f1f..35d8c63 100644
--- a/integrations/tensorflow/e2e/einsum_vector_test.py
+++ b/integrations/tensorflow/e2e/einsum_vector_test.py
@@ -14,8 +14,8 @@
# limitations under the License.
"""Test matrix ops via einsum"""
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
VECTOR_DIM = 16
diff --git a/integrations/tensorflow/e2e/fft_test.py b/integrations/tensorflow/e2e/fft_test.py
index 01bdfad..c16a772 100644
--- a/integrations/tensorflow/e2e/fft_test.py
+++ b/integrations/tensorflow/e2e/fft_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/fill_test.py b/integrations/tensorflow/e2e/fill_test.py
index 24f74ed..0c9f36d 100644
--- a/integrations/tensorflow/e2e/fill_test.py
+++ b/integrations/tensorflow/e2e/fill_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/gather_test.py b/integrations/tensorflow/e2e/gather_test.py
index fa10205..d27c67e 100644
--- a/integrations/tensorflow/e2e/gather_test.py
+++ b/integrations/tensorflow/e2e/gather_test.py
@@ -13,9 +13,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/image_resize_test.py b/integrations/tensorflow/e2e/image_resize_test.py
index 51d6c7d..ede5f66 100644
--- a/integrations/tensorflow/e2e/image_resize_test.py
+++ b/integrations/tensorflow/e2e/image_resize_test.py
@@ -13,9 +13,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_utils
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v1 as tf
diff --git a/integrations/tensorflow/e2e/keras/applications/applications_test.py b/integrations/tensorflow/e2e/keras/applications/applications_test.py
index 8c0e77c..50b6481 100644
--- a/integrations/tensorflow/e2e/keras/applications/applications_test.py
+++ b/integrations/tensorflow/e2e/keras/applications/applications_test.py
@@ -18,9 +18,9 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/keras/keyword_spotting_streaming_test.py b/integrations/tensorflow/e2e/keras/keyword_spotting_streaming_test.py
index e2da0c2..4cc92f2 100644
--- a/integrations/tensorflow/e2e/keras/keyword_spotting_streaming_test.py
+++ b/integrations/tensorflow/e2e/keras/keyword_spotting_streaming_test.py
@@ -18,8 +18,8 @@
from absl import app
from absl import flags
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
from kws_streaming.layers import modes
diff --git a/integrations/tensorflow/e2e/keras/layers/layers_test.py b/integrations/tensorflow/e2e/keras/layers/layers_test.py
index 1734d68..7d964c2 100644
--- a/integrations/tensorflow/e2e/keras/layers/layers_test.py
+++ b/integrations/tensorflow/e2e/keras/layers/layers_test.py
@@ -23,8 +23,8 @@
from absl import app
from absl import flags
from absl import logging
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/keras/train/classification_training_test.py b/integrations/tensorflow/e2e/keras/train/classification_training_test.py
index 2565906..335b282 100644
--- a/integrations/tensorflow/e2e/keras/train/classification_training_test.py
+++ b/integrations/tensorflow/e2e/keras/train/classification_training_test.py
@@ -18,8 +18,8 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/keras/train/regression_training_test.py b/integrations/tensorflow/e2e/keras/train/regression_training_test.py
index 9d9ac79..66f8bf0 100644
--- a/integrations/tensorflow/e2e/keras/train/regression_training_test.py
+++ b/integrations/tensorflow/e2e/keras/train/regression_training_test.py
@@ -18,8 +18,8 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/linspace_test.py b/integrations/tensorflow/e2e/linspace_test.py
index 8012918..a9708ca 100644
--- a/integrations/tensorflow/e2e/linspace_test.py
+++ b/integrations/tensorflow/e2e/linspace_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/mandelbrot_test.py b/integrations/tensorflow/e2e/mandelbrot_test.py
index 90c6fa1..bd400c3 100644
--- a/integrations/tensorflow/e2e/mandelbrot_test.py
+++ b/integrations/tensorflow/e2e/mandelbrot_test.py
@@ -13,7 +13,7 @@
# limitations under the License.
from absl import app
-from pyiree.tf.support import tf_test_utils
+from iree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/math/math_test.py b/integrations/tensorflow/e2e/math/math_test.py
index d98ffea..1ed7f74 100644
--- a/integrations/tensorflow/e2e/math/math_test.py
+++ b/integrations/tensorflow/e2e/math/math_test.py
@@ -18,9 +18,9 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/matrix_ops_dynamic_test.py b/integrations/tensorflow/e2e/matrix_ops_dynamic_test.py
index 0a3a766..10db11c 100644
--- a/integrations/tensorflow/e2e/matrix_ops_dynamic_test.py
+++ b/integrations/tensorflow/e2e/matrix_ops_dynamic_test.py
@@ -15,8 +15,8 @@
"""Test matrix ops."""
from absl import app
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/matrix_ops_static_test.py b/integrations/tensorflow/e2e/matrix_ops_static_test.py
index 67d2454..91d52e5 100644
--- a/integrations/tensorflow/e2e/matrix_ops_static_test.py
+++ b/integrations/tensorflow/e2e/matrix_ops_static_test.py
@@ -15,8 +15,8 @@
"""Test matrix ops."""
from absl import app
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
LEFT_DIM = 64
diff --git a/integrations/tensorflow/e2e/mobile_bert_squad_test.py b/integrations/tensorflow/e2e/mobile_bert_squad_test.py
index d0d8726..c94c397 100644
--- a/integrations/tensorflow/e2e/mobile_bert_squad_test.py
+++ b/integrations/tensorflow/e2e/mobile_bert_squad_test.py
@@ -23,8 +23,8 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
FLAGS = flags.FLAGS
diff --git a/integrations/tensorflow/e2e/quantization_dyn_test.py b/integrations/tensorflow/e2e/quantization_dyn_test.py
index d1c44ef..85c0e50 100644
--- a/integrations/tensorflow/e2e/quantization_dyn_test.py
+++ b/integrations/tensorflow/e2e/quantization_dyn_test.py
@@ -15,9 +15,9 @@
"""Tests for ops in the tf.math module."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/quantization_test.py b/integrations/tensorflow/e2e/quantization_test.py
index 2ccf8f8..edd8d27 100644
--- a/integrations/tensorflow/e2e/quantization_test.py
+++ b/integrations/tensorflow/e2e/quantization_test.py
@@ -15,9 +15,9 @@
"""Tests for ops in the tf.math module."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/range_test.py b/integrations/tensorflow/e2e/range_test.py
index ca988f6..0377ebb 100644
--- a/integrations/tensorflow/e2e/range_test.py
+++ b/integrations/tensorflow/e2e/range_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/resource_ops_test.py b/integrations/tensorflow/e2e/resource_ops_test.py
index 886649c..dbab65d 100644
--- a/integrations/tensorflow/e2e/resource_ops_test.py
+++ b/integrations/tensorflow/e2e/resource_ops_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/ring_buffer_test.py b/integrations/tensorflow/e2e/ring_buffer_test.py
index 18637ac..91d281d 100644
--- a/integrations/tensorflow/e2e/ring_buffer_test.py
+++ b/integrations/tensorflow/e2e/ring_buffer_test.py
@@ -14,8 +14,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
TIME_SIZE = 2
diff --git a/integrations/tensorflow/e2e/scatter_update_test.py b/integrations/tensorflow/e2e/scatter_update_test.py
index 9ed561c..da9819a 100644
--- a/integrations/tensorflow/e2e/scatter_update_test.py
+++ b/integrations/tensorflow/e2e/scatter_update_test.py
@@ -14,8 +14,8 @@
"""Test scatter update behavior for tensorflow."""
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/simple_arithmetic_test.py b/integrations/tensorflow/e2e/simple_arithmetic_test.py
index 8bd7b4e..8a27291 100644
--- a/integrations/tensorflow/e2e/simple_arithmetic_test.py
+++ b/integrations/tensorflow/e2e/simple_arithmetic_test.py
@@ -15,9 +15,9 @@
"""Several baseline e2e simple arithmetic tests."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/simple_stateful_test.py b/integrations/tensorflow/e2e/simple_stateful_test.py
index aa2cffa..b0ed6db 100644
--- a/integrations/tensorflow/e2e/simple_stateful_test.py
+++ b/integrations/tensorflow/e2e/simple_stateful_test.py
@@ -14,8 +14,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/sliding_window_test.py b/integrations/tensorflow/e2e/sliding_window_test.py
index 555cfa5..d6d9c37 100644
--- a/integrations/tensorflow/e2e/sliding_window_test.py
+++ b/integrations/tensorflow/e2e/sliding_window_test.py
@@ -13,8 +13,8 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
import tensorflow.compat.v2 as tf
TIME_SIZE = 3
diff --git a/integrations/tensorflow/e2e/slim_vision_models/slim_vision_model_test.py b/integrations/tensorflow/e2e/slim_vision_models/slim_vision_model_test.py
index 1d2fb5f..f6d58c5 100644
--- a/integrations/tensorflow/e2e/slim_vision_models/slim_vision_model_test.py
+++ b/integrations/tensorflow/e2e/slim_vision_models/slim_vision_model_test.py
@@ -18,9 +18,9 @@
from absl import app
from absl import flags
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow as tf
import tensorflow_hub as hub
diff --git a/integrations/tensorflow/e2e/space_to_batch_nd_test.py b/integrations/tensorflow/e2e/space_to_batch_nd_test.py
index 1e4c4cc..93f9ee7 100644
--- a/integrations/tensorflow/e2e/space_to_batch_nd_test.py
+++ b/integrations/tensorflow/e2e/space_to_batch_nd_test.py
@@ -15,9 +15,9 @@
"""Space To Batch ND tests."""
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/strings_test.py b/integrations/tensorflow/e2e/strings_test.py
index e8ef23d..09e9ec2 100644
--- a/integrations/tensorflow/e2e/strings_test.py
+++ b/integrations/tensorflow/e2e/strings_test.py
@@ -12,10 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from absl import app
-import numpy as np
-from pyiree.tf.support import tf_test_utils
import string
+
+from absl import app
+from iree.tf.support import tf_test_utils
+import numpy as np
import tensorflow.compat.v2 as tf
diff --git a/integrations/tensorflow/e2e/tensorlist_test.py b/integrations/tensorflow/e2e/tensorlist_test.py
index 104ede9..9de8ee4 100644
--- a/integrations/tensorflow/e2e/tensorlist_test.py
+++ b/integrations/tensorflow/e2e/tensorlist_test.py
@@ -13,9 +13,9 @@
# limitations under the License.
from absl import app
+from iree.tf.support import tf_test_utils
+from iree.tf.support import tf_utils
import numpy as np
-from pyiree.tf.support import tf_test_utils
-from pyiree.tf.support import tf_utils
import tensorflow.compat.v2 as tf
STATIC_SIZE = 20
