Use TF's py_test macro to make sure tests can run under pip builds.

PiperOrigin-RevId: 295072178
Change-Id: Ic2454e70e12e3bea76a2094273144e6b5ae57ffb

tensorflow/python/keras/distribute/BUILD

@@ -3,6 +3,7 @@
 #   related to dist-strat used by Keras..
 
 load("//tensorflow/core/platform/default:distribute.bzl", "distribute_py_test")
+load("//tensorflow:tensorflow.bzl", "py_test")
 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
 
 package(