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/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
