From 498c9045ff582bfc227f7c8769958c4394d537ca Mon Sep 17 00:00:00 2001
From: Scott Zhu <scottzhu@google.com>
Date: Fri, 6 Nov 2020 13:23:04 -0800
Subject: [PATCH] Cleanup the build file for keras package.

PiperOrigin-RevId: 341107268
Change-Id: I69706c7ce651f1fb207e97488f61d4cd78cda1a3
---
 tensorflow/python/keras/distribute/BUILD       |  8 ++++----
 tensorflow/python/keras/feature_column/BUILD   |  4 ++--
 tensorflow/python/keras/integration_test/BUILD |  3 ++-
 tensorflow/python/keras/mixed_precision/BUILD  | 10 +++++-----
 tensorflow/python/keras/premade/BUILD          |  6 +++---
 5 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
index f52b63defb4..87269568c26 100644
--- a/tensorflow/python/keras/distribute/BUILD
+++ b/tensorflow/python/keras/distribute/BUILD
@@ -3,8 +3,8 @@
 #   related to dist-strat used by Keras..
 
 load("//tensorflow/core/platform/default:distribute.bzl", "distribute_py_test")
-load("//tensorflow:tensorflow.bzl", "py_test", "tf_py_test")
 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")  # buildifier: disable=same-origin-load
 
 package(
     # TODO(scottzhu): Remove this deps when distribute test are converted to integration test.
@@ -747,7 +747,7 @@ cuda_py_test(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "multi_worker_callback_tf2_test",
     srcs = ["multi_worker_callback_tf2_test.py"],
     python_version = "PY3",
@@ -854,7 +854,7 @@ distribute_py_test(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "parameter_server_training_test",
     srcs = ["parameter_server_training_test.py"],
     python_version = "PY3",
@@ -894,7 +894,7 @@ py_library(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "distributed_file_utils_test",
     srcs = ["distributed_file_utils_test.py"],
     python_version = "PY3",
diff --git a/tensorflow/python/keras/feature_column/BUILD b/tensorflow/python/keras/feature_column/BUILD
index 89eaa6d1dd2..82661738a87 100644
--- a/tensorflow/python/keras/feature_column/BUILD
+++ b/tensorflow/python/keras/feature_column/BUILD
@@ -1,4 +1,4 @@
-load("//tensorflow:tensorflow.bzl", "py_test", "tf_py_test")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
 
 package(
     default_visibility = [
@@ -153,7 +153,7 @@ tf_py_test(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "sequence_feature_column_integration_test",
     srcs = ["sequence_feature_column_integration_test.py"],
     python_version = "PY3",
diff --git a/tensorflow/python/keras/integration_test/BUILD b/tensorflow/python/keras/integration_test/BUILD
index 164b731d56b..1190d7897f2 100644
--- a/tensorflow/python/keras/integration_test/BUILD
+++ b/tensorflow/python/keras/integration_test/BUILD
@@ -1,7 +1,8 @@
 # Description:
 #   Contains Keras integration tests that verify with other TF high level APIs.
 
-load("//tensorflow:tensorflow.bzl", "cuda_py_test", "tf_py_test")
+load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")  # buildifier: disable=same-origin-load
 load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_test")
 
 package(
diff --git a/tensorflow/python/keras/mixed_precision/BUILD b/tensorflow/python/keras/mixed_precision/BUILD
index d10ab72082f..ccf55c8b166 100644
--- a/tensorflow/python/keras/mixed_precision/BUILD
+++ b/tensorflow/python/keras/mixed_precision/BUILD
@@ -16,8 +16,8 @@
 # Description:
 #   Contains the Keras Mixed Precision API (TensorFlow version).
 
-load("//tensorflow:tensorflow.bzl", "py_test")
 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")  # buildifier: disable=same-origin-load
 
 package(
     default_visibility = [
@@ -61,7 +61,7 @@ py_library(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "policy_test",
     size = "medium",
     srcs = [
@@ -111,7 +111,7 @@ py_library(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "get_layer_policy_test",
     srcs = ["get_layer_policy_test.py"],
     srcs_version = "PY2AND3",
@@ -141,7 +141,7 @@ py_library(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "autocast_variable_test",
     size = "medium",
     srcs = ["autocast_variable_test.py"],
@@ -280,7 +280,7 @@ cuda_py_test(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "layer_correctness_test",
     size = "medium",
     srcs = ["layer_correctness_test.py"],
diff --git a/tensorflow/python/keras/premade/BUILD b/tensorflow/python/keras/premade/BUILD
index 9b3852e7ce2..b7570dad580 100644
--- a/tensorflow/python/keras/premade/BUILD
+++ b/tensorflow/python/keras/premade/BUILD
@@ -1,6 +1,6 @@
 # Description:
 #   Contains the Keras Premade Models (internal TensorFlow version).
-load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
 
 package(
     default_visibility = [
@@ -33,7 +33,7 @@ py_library(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "linear_test",
     size = "medium",
     srcs = ["linear_test.py"],
@@ -47,7 +47,7 @@ py_test(
     ],
 )
 
-py_test(
+tf_py_test(
     name = "wide_deep_test",
     size = "medium",
     srcs = ["wide_deep_test.py"],