Cleanup the build file for keras package.

PiperOrigin-RevId: 341107268
Change-Id: I69706c7ce651f1fb207e97488f61d4cd78cda1a3

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",

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",

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(

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"],

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"],