Move the module test wrt to Keras private API to keras test.

PiperOrigin-RevId: 316695477
Change-Id: I4201f62e97bb1acb768b980a51d26bf0c76853ab

tensorflow/python/keras/layers/BUILD

@@ -886,3 +886,15 @@ tf_py_test(
         "@absl_py//absl/testing:parameterized",
     ],
 )
+
+tf_py_test(
+    name = "layers_test",
+    size = "small",
+    srcs = ["layers_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":layers",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:tf2",
+    ],
+)

tensorflow/python/keras/layers/layers_test.py

@@ -0,0 +1,40 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# pylint: disable=g-classes-have-attributes
+"""Tests for layers.__init__."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python import tf2
+from tensorflow.python.keras import layers
+from tensorflow.python.platform import test
+
+
+class LayersTest(test.TestCase):
+
+  def test_keras_private_symbol(self):
+    normalization_parent = layers.Normalization.__module__.split('.')[-1]
+    if tf2.enabled():
+      self.assertEqual('normalization', normalization_parent)
+      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
+    else:
+      self.assertEqual('normalization_v1', normalization_parent)
+      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/tools/api/tests/module_test.py

@@ -79,18 +79,6 @@ class ModuleTest(test.TestCase):
       tf.compat.v1.summary.FileWriter
     # pylint: enable=pointless-statement
 
-  def testInternalKerasImport(self):
-    # pylint: disable=g-import-not-at-top
-    from tensorflow.python.keras import layers
-    normalization_parent = layers.Normalization.__module__.split('.')[-1]
-    if tf._major_api_version == 2:
-      self.assertEqual('normalization', normalization_parent)
-      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    else:
-      self.assertEqual('normalization_v1', normalization_parent)
-      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    # pylint: enable=g-import-not-at-top
-
 
 if __name__ == '__main__':
   test.main()