Replace tensorflow.python.keras with keras. tensorflow.python.keras is an old copy and is deprecated.

PiperOrigin-RevId: 490106692

Author: divyashreepathihalli

tensorflow_model_optimization/python/core/quantization/keras/quant_ops_test.py

@@ -18,19 +18,16 @@
 from __future__ import division
 from __future__ import print_function
 
-from absl.testing import parameterized
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import compat
 from tensorflow_model_optimization.python.core.quantization.keras import quant_ops
 
 _SYMMETRIC_RANGE_RATIO = 0.9921875  # 127 / 128
 
 
-@keras_parameterized.run_all_keras_modes
-class QuantOpsTest(tf.test.TestCase, parameterized.TestCase):
+class QuantOpsTest(tf.test.TestCase):
 
   def testAllValuesQuantiize_TrainingAssign(self):
     min_value, max_value = self._GetMinMaxValues(

tensorflow_model_optimization/python/core/quantization/keras/quantize_aware_activation_test.py

@@ -23,7 +23,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 
@@ -37,7 +36,8 @@
 MovingAverageQuantizer = quantizers.MovingAverageQuantizer
 
 
-@keras_parameterized.run_all_keras_modes
+@tf.__internal__.distribute.combinations.generate(
+    tf.__internal__.test.combinations.combine(mode=['graph', 'eager']))
 class QuantizeAwareQuantizationTest(tf.test.TestCase, parameterized.TestCase):
 
   def setUp(self):

tensorflow_model_optimization/python/core/quantization/keras/quantize_functional_test.py

@@ -25,7 +25,6 @@
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import compat
 from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
 from tensorflow_model_optimization.python.core.quantization.keras import quantize
@@ -34,8 +33,9 @@
 layers = tf.keras.layers
 
 
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
-class QuantizeFunctionalTest(tf.test.TestCase, parameterized.TestCase):
+@tf.__internal__.distribute.combinations.generate(
+    tf.__internal__.test.combinations.combine(mode=['graph', 'eager']))
+class QuantizeFunctionalTest(tf.test.TestCase):
 
   # TODO(pulkitb): Parameterize test and include functional mnist, and
   # other RNN models.

tensorflow_model_optimization/python/core/quantization/keras/quantize_integration_test.py

@@ -26,7 +26,6 @@
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
-from tensorflow.python.keras import keras_parameterized
 
 from tensorflow_model_optimization.python.core.keras import compat
 from tensorflow_model_optimization.python.core.keras import test_utils
@@ -44,7 +43,8 @@
 # TODO(tfmot): enable for v1. Currently fails because the decorator
 # on graph mode wraps everything in a graph, which is not compatible
 # with the TFLite converter's call to clear_session().
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
+@tf.__internal__.distribute.combinations.generate(
+    tf.__internal__.test.combinations.combine(mode=['graph', 'eager']))
 class QuantizeIntegrationTest(tf.test.TestCase, parameterized.TestCase):
 
   def _batch(self, dims, batch_size):

tensorflow_model_optimization/python/core/quantization/keras/quantize_models_test.py

@@ -26,12 +26,10 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.quantization.keras import quantize
 from tensorflow_model_optimization.python.core.quantization.keras import utils
 
 
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
 class QuantizeModelsTest(tf.test.TestCase, parameterized.TestCase):
 
   # Derived using

tensorflow_model_optimization/python/core/quantization/keras/quantizers_test.py

@@ -23,15 +23,13 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized
 from tensorflow_model_optimization.python.core.keras import compat
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 
 deserialize_keras_object = tf.keras.utils.deserialize_keras_object
 serialize_keras_object = tf.keras.utils.serialize_keras_object
 
 
-@keras_parameterized.run_all_keras_modes
 @parameterized.parameters(
     quantizers.LastValueQuantizer,
     quantizers.MovingAverageQuantizer,