Move keras related module test to keras/integration_test.

PiperOrigin-RevId: 306685976
Change-Id: I3d674da22e5a919048298773ebe32c60338e5fba

tensorflow/python/keras/integration_test/BUILD

@@ -51,3 +51,12 @@ tf_py_test(
         "//tensorflow/python:extra_py_tests_deps",
     ],
 )
+
+tf_py_test(
+    name = "module_test",
+    srcs = ["module_test.py"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python:extra_py_tests_deps",
+    ],
+)

tensorflow/python/keras/integration_test/module_test.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class ModuleTest(tf.test.TestCase):
+
+  def test_module_discover_layer_variable(self):
+    m = tf.Module()
+    m.a = tf.keras.layers.Dense(1)
+    m.b = tf.keras.layers.Dense(2)
+
+    # The weights of the layer has not been created yet.
+    self.assertEmpty(m.variables)
+    self.assertLen(m.submodules, 2)
+
+    inputs = tf.keras.layers.Input((1,))
+    m.a(inputs)
+    m.b(inputs)
+
+    variable_list = m.variables
+    self.assertLen(variable_list, 4)
+    self.assertIs(variable_list[0], m.a.kernel)
+    self.assertIs(variable_list[1], m.a.bias)
+    self.assertIs(variable_list[2], m.b.kernel)
+    self.assertIs(variable_list[3], m.b.bias)
+
+  def test_model_discover_submodule(self):
+    m = tf.keras.models.Sequential(
+        layers=[tf.keras.layers.Dense(1), tf.keras.layers.Dense(2)])
+
+    self.assertEqual(m.submodules, (m.layers[0], m.layers[1]))
+    m(tf.keras.layers.Input((1,)))
+    self.assertLen(m.variables, 4)
+
+  def test_model_wrapped_in_module_discovers_submodules(self):
+    linear = tf.keras.models.Sequential(
+        [tf.keras.layers.Dense(units=1, input_shape=[1])])
+    linear.compile(optimizer="sgd", loss="mean_squared_error")
+    m = tf.Module()
+    m.l = linear
+    self.assertNotEmpty(m.submodules)
+    self.assertLen(m.variables, 2)
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensorflow/python/module/module_test.py

@@ -32,8 +32,6 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
-from tensorflow.python.keras import layers
-from tensorflow.python.keras import models
 from tensorflow.python.module import module
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
@@ -505,42 +503,6 @@ class FlattenTest(parameterized.TestCase, test_util.TensorFlowTestCase):
                       ("decoder", "w", 0, 0, "k"): mod.decoder.w[0][0]["k"],
                       ("decoder", "w", 0, 1, "k"): mod.decoder.w[0][1]["k"]},)
 
-  def test_module_discover_layer_variable(self):
-    m = module.Module()
-    m.a = layers.Dense(1)
-    m.b = layers.Dense(2)
-
-    # The weights of the layer has not been created yet.
-    self.assertEmpty(m.variables)
-    self.assertLen(m.submodules, 2)
-
-    inputs = layers.Input((1,))
-    m.a(inputs)
-    m.b(inputs)
-
-    variable_list = m.variables
-    self.assertLen(variable_list, 4)
-    self.assertIs(variable_list[0], m.a.kernel)
-    self.assertIs(variable_list[1], m.a.bias)
-    self.assertIs(variable_list[2], m.b.kernel)
-    self.assertIs(variable_list[3], m.b.bias)
-
-  def test_model_discover_submodule(self):
-    m = models.Sequential(layers=[layers.Dense(1),
-                                  layers.Dense(2)])
-
-    self.assertEqual(m.submodules, (m.layers[0], m.layers[1]))
-    m(layers.Input((1,)))
-    self.assertLen(m.variables, 4)
-
-  def test_model_wrapped_in_module_discovers_submodules(self):
-    linear = models.Sequential([layers.Dense(units=1, input_shape=[1])])
-    linear.compile(optimizer="sgd", loss="mean_squared_error")
-    m = module.Module()
-    m.l = linear
-    self.assertNotEmpty(m.submodules)
-    self.assertLen(m.variables, 2)
-
   def test_raises_error_with_path(self):
     if six.PY2:
       class NonOrderable(object):