Remove private Keras imports.

PiperOrigin-RevId: 564511437

tensorflow_text/python/keras/layers/todense_test.py

@@ -19,8 +19,6 @@
 from __future__ import print_function
 
 from absl.testing import parameterized
-from keras.testing_infra import test_combinations
-from keras.testing_infra import test_utils as keras_test_utils
 import numpy as np
 import tensorflow as tf
 
@@ -52,26 +50,40 @@ def get_input_dataset(in_data, out_data=None):
       (in_data, out_data)).batch(batch_size)
 
 
-@test_combinations.run_with_all_model_types
-@test_combinations.run_all_keras_modes
-class RaggedTensorsToDenseLayerTest(test_combinations.TestCase):
+def get_model_from_layers(
+    layers,
+    input_shape,
+    input_sparse=False,
+    input_ragged=False,
+    input_dtype=None):
+  layers = [
+      tf.keras.Input(
+          shape=input_shape,
+          dtype=input_dtype,
+          sparse=input_sparse,
+          ragged=input_ragged,
+      )
+  ] + layers
+  return tf.keras.models.Sequential(layers)
+
+
+class RaggedTensorsToDenseLayerTest(tf.test.TestCase, parameterized.TestCase):
 
   def SKIP_test_ragged_input_default_padding(self):
     input_data = get_input_dataset(
         tf.ragged.constant([[1, 2, 3, 4, 5], [2, 3]]))
     expected_output = np.array([[1, 2, 3, 4, 5], [2, 3, 0, 0, 0]])
 
     layers = [ToDense(), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_ragged=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)
 
@@ -84,16 +96,15 @@ def SKIP_test_ragged_input_with_padding(self):
                                  [3., -1., -1., -1., -1.]]])
 
     layers = [ToDense(pad_value=-1), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None, None),
         input_ragged=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)
 
@@ -113,16 +124,15 @@ def test_ragged_input_shape(self):
     expected_output = np.array([[1, 2, 3, 4, 5, 0, 0], [2, 3, 0, 0, 0, 0, 0]])
 
     layers = [ToDense(shape=[2, 7]), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_ragged=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)
 
@@ -132,7 +142,7 @@ def test_ragged_input_shape(self):
           tf.compat.v1.keras.layers.LSTM, tf.keras.layers.GRU,
           tf.keras.layers.LSTM
       ]))
-  def SKIP_test_ragged_input_RNN_layer(self, layer):
+  def SKIP_test_ragged_input_RNN_layer(self, layer):  # pylint: disable=invalid-name
     input_data = get_input_dataset(
         tf.ragged.constant([[1, 2, 3, 4, 5], [5, 6]]))
 
@@ -143,24 +153,21 @@ def SKIP_test_ragged_input_RNN_layer(self, layer):
         tf.keras.layers.Dense(3, activation="softmax"),
         tf.keras.layers.Dense(1, activation="sigmoid")
     ]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_ragged=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="rmsprop",
         loss="binary_crossentropy",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
 
     output = model.predict(input_data)
     self.assertAllEqual(np.zeros((2, 1)).shape, output.shape)
 
 
-@test_combinations.run_with_all_model_types
-@test_combinations.run_all_keras_modes
-class SparseTensorsToDenseLayerTest(test_combinations.TestCase):
+class SparseTensorsToDenseLayerTest(tf.test.TestCase):
 
   def SKIP_test_sparse_input_default_padding(self):
     input_data = get_input_dataset(
@@ -171,16 +178,15 @@ def SKIP_test_sparse_input_default_padding(self):
                                 [0., 0., 0., 0.]])
 
     layers = [ToDense(), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_sparse=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)
 
@@ -193,16 +199,15 @@ def SKIP_test_sparse_input_with_padding(self):
                                 [-1., -1., -1., -1.]])
 
     layers = [ToDense(pad_value=-1, trainable=False), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_sparse=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)
 
@@ -227,16 +232,15 @@ def test_sparse_input_shape(self):
                                 [0., 0., 0., 0.]])
 
     layers = [ToDense(shape=[3, 4]), Final()]
-    model = keras_test_utils.get_model_from_layers(
+    model = get_model_from_layers(
         layers,
         input_shape=(None,),
         input_sparse=True,
         input_dtype=tf.dtypes.int32)
     model.compile(
         optimizer="sgd",
         loss="mse",
-        metrics=["accuracy"],
-        run_eagerly=keras_test_utils.should_run_eagerly())
+        metrics=["accuracy"])
     output = model.predict(input_data)
     self.assertAllEqual(output, expected_output)