Enable presubmit testing for the mlir bridge.
These tests will run with both the new bridge and the old bridge.
PiperOrigin-RevId: 325061701
Change-Id: Id8802c48ac1693aa5b801207bc3f87cd303f03c4
diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index c1b0ee6..7965be1 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -851,6 +851,7 @@
distribute_py_test(
name = "strategy_combinations_test",
srcs = ["strategy_combinations_test.py"],
+ disable_mlir_bridge = False,
python_version = "PY3",
deps = [
":combinations",
@@ -905,6 +906,7 @@
distribute_py_test(
name = "checkpointing_test",
srcs = ["checkpointing_test.py"],
+ disable_mlir_bridge = False,
main = "checkpointing_test.py",
tags = [
"multi_and_single_gpu",
@@ -920,6 +922,7 @@
distribute_py_test(
name = "input_lib_test",
srcs = ["input_lib_test.py"],
+ disable_mlir_bridge = False,
main = "input_lib_test.py",
shard_count = 10,
tags = [
@@ -1116,6 +1119,7 @@
name = "values_test",
size = "medium",
srcs = ["values_test.py"],
+ disable_mlir_bridge = False,
main = "values_test.py",
shard_count = 5,
tags = [
@@ -1173,6 +1177,7 @@
distribute_py_test(
name = "distribute_utils_test",
srcs = ["distribute_utils_test.py"],
+ disable_mlir_bridge = False,
main = "distribute_utils_test.py",
tags = [
"multi_and_single_gpu",
@@ -1238,6 +1243,7 @@
name = "ps_values_test",
size = "medium",
srcs = ["ps_values_test.py"],
+ disable_mlir_bridge = False,
main = "ps_values_test.py",
tags = [
"multi_and_single_gpu",
@@ -1257,6 +1263,7 @@
distribute_py_test(
name = "moving_averages_test",
srcs = ["moving_averages_test.py"],
+ disable_mlir_bridge = False,
main = "moving_averages_test.py",
deps = [
":combinations",
@@ -1274,6 +1281,7 @@
distribute_py_test(
name = "custom_training_loop_gradient_test",
srcs = ["custom_training_loop_gradient_test.py"],
+ disable_mlir_bridge = False,
main = "custom_training_loop_gradient_test.py",
tags = [
"multi_and_single_gpu",
@@ -1472,6 +1480,7 @@
distribute_py_test(
name = "zero_batch_test",
srcs = ["zero_batch_test.py"],
+ disable_mlir_bridge = False,
main = "zero_batch_test.py",
deps = [
":combinations",
@@ -1671,6 +1680,7 @@
distribute_py_test(
name = "strategy_common_test",
srcs = ["strategy_common_test.py"],
+ disable_mlir_bridge = False,
python_version = "PY3",
shard_count = 2,
tags = [
@@ -1703,6 +1713,7 @@
distribute_py_test(
name = "tf_function_test",
srcs = ["tf_function_test.py"],
+ disable_mlir_bridge = False,
main = "tf_function_test.py",
tags = [
"multi_and_single_gpu",
diff --git a/tensorflow/python/distribute/integration_test/BUILD b/tensorflow/python/distribute/integration_test/BUILD
index d997e64..156699f 100644
--- a/tensorflow/python/distribute/integration_test/BUILD
+++ b/tensorflow/python/distribute/integration_test/BUILD
@@ -8,6 +8,7 @@
distribute_py_test(
name = "saved_model_test",
srcs = ["saved_model_test.py"],
+ disable_mlir_bridge = False,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python/distribute:combinations",
diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
index 56a6a9d..2a0421c 100644
--- a/tensorflow/python/keras/distribute/BUILD
+++ b/tensorflow/python/keras/distribute/BUILD
@@ -144,6 +144,7 @@
distribute_py_test(
name = "checkpointing_test",
srcs = ["checkpointing_test.py"],
+ disable_mlir_bridge = False,
main = "checkpointing_test.py",
tags = [
"multi_and_single_gpu",
@@ -221,6 +222,7 @@
distribute_py_test(
name = "custom_training_loop_metrics_test",
srcs = ["custom_training_loop_metrics_test.py"],
+ disable_mlir_bridge = False,
main = "custom_training_loop_metrics_test.py",
tags = [
"multi_and_single_gpu",
@@ -267,6 +269,7 @@
distribute_py_test(
name = "custom_training_loop_optimizer_test",
srcs = ["custom_training_loop_optimizer_test.py"],
+ disable_mlir_bridge = False,
main = "custom_training_loop_optimizer_test.py",
tags = [
"multi_and_single_gpu",
@@ -311,6 +314,7 @@
distribute_py_test(
name = "keras_premade_models_test",
srcs = ["keras_premade_models_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "keras_premade_models_test.py",
shard_count = 4,
@@ -347,6 +351,7 @@
distribute_py_test(
name = "distributed_training_utils_test",
srcs = ["distributed_training_utils_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "distributed_training_utils_test.py",
deps = [
@@ -389,6 +394,7 @@
name = "keras_dnn_correctness_test",
size = "medium",
srcs = ["keras_dnn_correctness_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "keras_dnn_correctness_test.py",
# Shard count is set to an odd number to distribute tasks across
@@ -426,6 +432,7 @@
name = "keras_image_model_correctness_test",
size = "medium",
srcs = ["keras_image_model_correctness_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "keras_image_model_correctness_test.py",
shard_count = 16,
@@ -444,6 +451,7 @@
distribute_py_test(
name = "keras_metrics_test",
srcs = ["keras_metrics_test.py"],
+ disable_mlir_bridge = False,
main = "keras_metrics_test.py",
tags = [
"multi_and_single_gpu",
@@ -484,6 +492,7 @@
name = "keras_save_load_test",
size = "medium",
srcs = ["keras_save_load_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "keras_save_load_test.py",
shard_count = 7,
@@ -500,6 +509,7 @@
name = "keras_stateful_lstm_model_correctness_test",
size = "medium",
srcs = ["keras_stateful_lstm_model_correctness_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "keras_stateful_lstm_model_correctness_test.py",
shard_count = 4,
@@ -749,6 +759,7 @@
name = "saved_model_save_load_test",
size = "medium",
srcs = ["saved_model_save_load_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "saved_model_save_load_test.py",
shard_count = 7,
@@ -766,6 +777,7 @@
name = "saved_model_mixed_api_test",
size = "medium",
srcs = ["saved_model_mixed_api_test.py"],
+ disable_mlir_bridge = False,
full_precision = True,
main = "saved_model_mixed_api_test.py",
shard_count = 7,
diff --git a/tensorflow/python/keras/integration_test/BUILD b/tensorflow/python/keras/integration_test/BUILD
index b23dcc5..20e1a88 100644
--- a/tensorflow/python/keras/integration_test/BUILD
+++ b/tensorflow/python/keras/integration_test/BUILD
@@ -97,6 +97,7 @@
name = "tpu_strategy_test",
srcs = ["tpu_strategy_test.py"],
disable_experimental = True,
+ disable_mlir_bridge = False,
python_version = "PY3",
tags = ["no_oss"],
deps = [
diff --git a/tensorflow/python/keras/layers/preprocessing/BUILD b/tensorflow/python/keras/layers/preprocessing/BUILD
index 1fa6deb..3e6624b 100644
--- a/tensorflow/python/keras/layers/preprocessing/BUILD
+++ b/tensorflow/python/keras/layers/preprocessing/BUILD
@@ -326,6 +326,7 @@
distribute_py_test(
name = "category_encoding_distribution_test",
srcs = ["category_encoding_distribution_test.py"],
+ disable_mlir_bridge = False,
main = "category_encoding_distribution_test.py",
python_version = "PY3",
tags = [
@@ -423,6 +424,7 @@
tpu_py_test(
name = "hashing_distribution_test",
srcs = ["hashing_distribution_test.py"],
+ disable_mlir_bridge = False,
main = "hashing_distribution_test.py",
python_version = "PY3",
tags = ["multi_and_single_gpu"],
@@ -453,6 +455,7 @@
tpu_py_test(
name = "index_lookup_distribution_test",
srcs = ["index_lookup_distribution_test.py"],
+ disable_mlir_bridge = False,
main = "index_lookup_distribution_test.py",
python_version = "PY3",
tags = ["no_oss"],
diff --git a/tensorflow/python/tpu/BUILD b/tensorflow/python/tpu/BUILD
index e1e71e6..9cd75d1 100644
--- a/tensorflow/python/tpu/BUILD
+++ b/tensorflow/python/tpu/BUILD
@@ -474,6 +474,7 @@
"tpu_embedding_v2_test.py",
],
disable_experimental = True,
+ disable_mlir_bridge = False,
python_version = "PY3",
shard_count = 4,
srcs_version = "PY2AND3",
@@ -505,6 +506,7 @@
"tpu_embedding_v2_correctness_test.py",
],
disable_experimental = True,
+ disable_mlir_bridge = False,
python_version = "PY3",
shard_count = 4,
srcs_version = "PY2AND3",
