Update all the MG nlp tests to remove the keras_parameterized annotation.

Keras parameterized will run 3 modes:
1. v1_session is ignored when eager execution is enabled (default on in tf2).
2. v2_functional and v2_eager will run same exactly same code when the test doesn't use model.compile(run_eagerly=testing_util.should_run_eagerly()).

The removal of the annotation will eliminate an import to keras private API, and also reduce 50% of test executed.

PiperOrigin-RevId: 494249251

official/nlp/modeling/layers/attention_test.py

@@ -17,7 +17,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import attention
 
 
@@ -32,8 +31,7 @@ def _create_cache(batch_size, init_decode_length, num_heads, head_size):
   }
 
 
-@keras_parameterized.run_all_keras_modes
-class CachedAttentionTest(keras_parameterized.TestCase):
+class CachedAttentionTest(tf.test.TestCase):
 
   def test_masked_attention(self):
     """Test with a mask tensor."""

official/nlp/modeling/layers/block_diag_feedforward_test.py

@@ -18,14 +18,12 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import block_diag_feedforward
 
 
 # This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
 # guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BlockDiagFeedforwardTest(keras_parameterized.TestCase):
+class BlockDiagFeedforwardTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(BlockDiagFeedforwardTest, self).tearDown()

official/nlp/modeling/layers/gated_feedforward_test.py

@@ -18,14 +18,10 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import gated_feedforward
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class GatedFeedforwardTest(keras_parameterized.TestCase):
+class GatedFeedforwardTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(GatedFeedforwardTest, self).tearDown()

official/nlp/modeling/layers/masked_lm_test.py

@@ -17,16 +17,11 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
-
 from official.nlp.modeling.layers import masked_lm
 from official.nlp.modeling.networks import bert_encoder
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class MaskedLMTest(keras_parameterized.TestCase):
+class MaskedLMTest(tf.test.TestCase):
 
   def create_layer(self,
                    vocab_size,

official/nlp/modeling/layers/masked_softmax_test.py

@@ -17,14 +17,10 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import masked_softmax
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class MaskedSoftmaxLayerTest(keras_parameterized.TestCase):
+class MaskedSoftmaxLayerTest(tf.test.TestCase):
 
   def test_non_masked_softmax(self):
     test_layer = masked_softmax.MaskedSoftmax()

official/nlp/modeling/layers/mat_mul_with_margin_test.py

@@ -16,11 +16,10 @@
 
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import mat_mul_with_margin
 
 
-class MatMulWithMarginTest(keras_parameterized.TestCase):
+class MatMulWithMarginTest(tf.test.TestCase):
 
   def test_layer_invocation(self):
     """Validate that the Keras object can be created and invoked."""

official/nlp/modeling/layers/on_device_embedding_test.py

@@ -17,14 +17,10 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import on_device_embedding
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class OnDeviceEmbeddingTest(keras_parameterized.TestCase):
+class OnDeviceEmbeddingTest(tf.test.TestCase):
 
   def test_layer_creation(self):
     vocab_size = 31

official/nlp/modeling/layers/position_embedding_test.py

@@ -18,14 +18,10 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import position_embedding
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class PositionEmbeddingLayerTest(keras_parameterized.TestCase):
+class PositionEmbeddingLayerTest(tf.test.TestCase):
 
   def test_static_layer_output_shape(self):
     # Create a 3-dimensional input (the first dimension is implicit).
@@ -129,10 +125,7 @@ class PositionEmbeddingLayerTest(keras_parameterized.TestCase):
     self.assertAllEqual([1, input_length, width], output_data.shape)
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class RelativePositionEmbeddingLayerTest(keras_parameterized.TestCase):
+class RelativePositionEmbeddingLayerTest(tf.test.TestCase):
 
   def test_relative_tensor_input(self):
     hidden_size = 8
@@ -164,8 +157,7 @@ class RelativePositionEmbeddingLayerTest(keras_parameterized.TestCase):
     self.assertAllEqual(output_tensor, expected_output_tensor)
 
 
-@keras_parameterized.run_all_keras_modes
-class RelativePositionBiasTest(keras_parameterized.TestCase):
+class RelativePositionBiasTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(("bidirectional", True),
                                   ("unidirectional", False))

official/nlp/modeling/layers/relative_attention_test.py

@@ -14,11 +14,11 @@
 
 """Tests for the attention layer."""
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
 from tensorflow.python.distribute import combinations
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import relative_attention
 
 
@@ -111,8 +111,7 @@ def _create_mock_attention_data(
   return data
 
 
-@keras_parameterized.run_all_keras_modes
-class MultiHeadRelativeAttentionTest(keras_parameterized.TestCase):
+class MultiHeadRelativeAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       value_dim=[32, 64],
@@ -147,8 +146,7 @@ class MultiHeadRelativeAttentionTest(keras_parameterized.TestCase):
     self.assertEqual(output.shape, [batch_size, seq_length, key_dim])
 
 
-@keras_parameterized.run_all_keras_modes
-class TwoStreamRelativeAttentionTest(keras_parameterized.TestCase):
+class TwoStreamRelativeAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       num_predictions=[2, 10],

official/nlp/modeling/layers/rezero_transformer_test.py

@@ -18,14 +18,10 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import rezero_transformer
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class TransformerWithReZeroLayerTest(keras_parameterized.TestCase):
+class TransformerWithReZeroLayerTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(TransformerWithReZeroLayerTest, self).tearDown()

official/nlp/modeling/layers/talking_heads_attention_test.py

@@ -18,15 +18,11 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import talking_heads_attention
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
 # This test is revised base on attention.MultiHeadAttentionTest.
-@keras_parameterized.run_all_keras_modes
-class TalkingHeadsAttentionTest(keras_parameterized.TestCase):
+class TalkingHeadsAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("key_value_same_proj", None, None, [40, 80]),

official/nlp/modeling/layers/tn_transformer_test.py

@@ -18,15 +18,11 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers.tn_transformer_expand_condense import TNTransformerExpandCondense
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
 @parameterized.named_parameters(('tn', TNTransformerExpandCondense))
-class TransformerLayerTest(keras_parameterized.TestCase):
+class TransformerLayerTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(TransformerLayerTest, self).tearDown()

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -18,13 +18,12 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers.transformer_encoder_block import TransformerEncoderBlock
 
 
-@keras_parameterized.run_all_keras_modes
 @parameterized.named_parameters(('base', TransformerEncoderBlock))
-class TransformerEncoderBlockLayerTest(keras_parameterized.TestCase):
+class TransformerEncoderBlockLayerTest(
+    tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(TransformerEncoderBlockLayerTest, self).tearDown()
@@ -258,9 +257,8 @@ class TransformerEncoderBlockLayerTest(keras_parameterized.TestCase):
     self.assertEqual(output.shape, q_tensor.shape)
 
 
-@keras_parameterized.run_all_keras_modes
-class TransformerEncoderBlockLayerTestWithoutParams(keras_parameterized.TestCase
-                                                   ):
+class TransformerEncoderBlockLayerTestWithoutParams(
+    tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(TransformerEncoderBlockLayerTestWithoutParams, self).tearDown()
@@ -407,8 +405,7 @@ class TransformerEncoderBlockLayerTestWithoutParams(keras_parameterized.TestCase
                      test_layer._attention_layer.get_config()['value_dim'])
 
 
-@keras_parameterized.run_all_keras_modes
-class TransformerArgumentTest(keras_parameterized.TestCase):
+class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
 
   def test_use_bias_norm_first(self):
     num_attention_heads = 2
@@ -681,7 +678,7 @@ class TransformerArgumentTest(keras_parameterized.TestCase):
 
     if return_attention_scores:
       self.assertIsInstance(output_tensor, tuple)
-      self.assertEqual(len(output_tensor), 2)
+      self.assertLen(output_tensor, 2)
       # First is the standard output.
       self.assertEqual(output_tensor[0].shape.as_list(),
                        expected_layer_output_shape)

official/nlp/modeling/layers/transformer_scaffold_test.py

@@ -17,7 +17,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import attention
 from official.nlp.modeling.layers import transformer_scaffold
 
@@ -76,10 +75,7 @@ class ValidatedFeedforwardLayer(tf.keras.layers.Layer):
     return config
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class TransformerLayerTest(keras_parameterized.TestCase):
+class TransformerLayerTest(tf.test.TestCase):
 
   def tearDown(self):
     super(TransformerLayerTest, self).tearDown()

official/nlp/modeling/layers/transformer_test.py

@@ -16,7 +16,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.layers import transformer
 
 
@@ -31,8 +30,7 @@ def _create_cache(batch_size, init_decode_length, num_heads, head_size):
   }
 
 
-@keras_parameterized.run_all_keras_modes
-class TransformerDecoderBlockTest(keras_parameterized.TestCase):
+class TransformerDecoderBlockTest(tf.test.TestCase):
 
   def test_decoder_block_with_cache(self):
     num_attention_heads = 2

official/nlp/modeling/layers/transformer_xl_test.py

@@ -14,11 +14,11 @@
 
 """Tests for Transformer XL."""
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
 from tensorflow.python.distribute import combinations
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 
 from official.nlp.modeling.layers import transformer_xl
 
@@ -115,8 +115,7 @@ def create_mock_transformer_xl_data(
   return data
 
 
-@keras_parameterized.run_all_keras_modes
-class TransformerXLBlockTest(keras_parameterized.TestCase):
+class TransformerXLBlockTest(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       memory_length=[0, 4],
@@ -186,8 +185,7 @@ class TransformerXLBlockTest(keras_parameterized.TestCase):
     self.assertEqual(transformer_xl_block_config, new_block.get_config())
 
 
-@keras_parameterized.run_all_keras_modes
-class TransformerXLTest(keras_parameterized.TestCase):
+class TransformerXLTest(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       two_stream=[True, False],
@@ -246,7 +244,7 @@ class TransformerXLTest(keras_parameterized.TestCase):
     else:
       self.assertEqual(attention_output.shape,
                        [batch_size, seq_length, hidden_size])
-    self.assertEqual(len(cached_memory_states), num_layers)
+    self.assertLen(cached_memory_states, num_layers)
 
   def test_get_config(self):
     transformer_xl_layer = transformer_xl.TransformerXL(

official/nlp/modeling/losses/weighted_sparse_categorical_crossentropy_test.py

@@ -17,14 +17,12 @@ import numpy as np
 
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import layers
 from official.nlp.modeling import networks
 from official.nlp.modeling.losses import weighted_sparse_categorical_crossentropy
 
 
-@keras_parameterized.run_all_keras_modes
-class ClassificationLossTest(keras_parameterized.TestCase):
+class ClassificationLossTest(tf.test.TestCase):
 
   def create_lm_model(self,
                       vocab_size,

official/nlp/modeling/models/bert_classifier_test.py

@@ -17,16 +17,12 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import layers
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import bert_classifier
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertClassifierTest(keras_parameterized.TestCase):
+class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(('single_cls', 1, False), ('3_cls', 3, False),
                                   ('3_cls_dictoutputs', 3, True))

official/nlp/modeling/models/bert_pretrainer_test.py

@@ -18,16 +18,12 @@ import itertools
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import layers
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import bert_pretrainer
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertPretrainerTest(keras_parameterized.TestCase):
+class BertPretrainerTest(tf.test.TestCase, parameterized.TestCase):
 
   def test_bert_pretrainer(self):
     """Validate that the Keras object can be created."""
@@ -109,7 +105,7 @@ class BertPretrainerTest(keras_parameterized.TestCase):
                         new_bert_trainer_model.get_config())
 
 
-class BertPretrainerV2Test(keras_parameterized.TestCase):
+class BertPretrainerV2Test(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(itertools.product(
       (False, True),
@@ -121,6 +117,7 @@ class BertPretrainerV2Test(keras_parameterized.TestCase):
                              use_customized_masked_lm, has_masked_lm_positions):
     """Validate that the Keras object can be created."""
     # Build a transformer network to use within the BERT trainer.
+    del dict_outputs, return_all_encoder_outputs
     vocab_size = 100
     sequence_length = 512
     hidden_size = 48

official/nlp/modeling/models/bert_span_labeler_test.py

@@ -17,15 +17,11 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import bert_span_labeler
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertSpanLabelerTest(keras_parameterized.TestCase):
+class BertSpanLabelerTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(True, False)
   def test_bert_trainer(self, dict_outputs):
@@ -48,7 +44,7 @@ class BertSpanLabelerTest(keras_parameterized.TestCase):
     cls_outs = bert_trainer_model([word_ids, mask, type_ids])
 
     # Validate that there are 2 outputs are of the expected shape.
-    self.assertEqual(2, len(cls_outs))
+    self.assertLen(cls_outs, 2)
     expected_shape = [None, sequence_length]
     for out in cls_outs:
       self.assertAllEqual(expected_shape, out.shape.as_list())

official/nlp/modeling/models/bert_token_classifier_test.py

@@ -17,15 +17,11 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import bert_token_classifier
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertTokenClassifierTest(keras_parameterized.TestCase):
+class BertTokenClassifierTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters((True, True), (False, False))
   def test_bert_trainer(self, dict_outputs, output_encoder_outputs):

official/nlp/modeling/models/dual_encoder_test.py

@@ -17,15 +17,11 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import dual_encoder
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class DualEncoderTest(keras_parameterized.TestCase):
+class DualEncoderTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters((192, 'logits'), (768, 'predictions'))
   def test_dual_encoder(self, hidden_size, output):
@@ -72,6 +68,7 @@ class DualEncoderTest(keras_parameterized.TestCase):
   def test_dual_encoder_tensor_call(self, hidden_size, output):
     """Validate that the Keras object can be invoked."""
     # Build a transformer network to use within the dual encoder model.
+    del hidden_size
     sequence_length = 2
     test_network = networks.BertEncoder(vocab_size=100, num_layers=2)
 

official/nlp/modeling/models/electra_pretrainer_test.py

@@ -16,15 +16,11 @@
 
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import electra_pretrainer
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class ElectraPretrainerTest(keras_parameterized.TestCase):
+class ElectraPretrainerTest(tf.test.TestCase):
 
   def test_electra_pretrainer(self):
     """Validate that the Keras object can be created."""

official/nlp/modeling/models/xlnet_test.py

@@ -19,7 +19,6 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling import networks
 from official.nlp.modeling.models import xlnet
 
@@ -44,10 +43,7 @@ def _get_xlnet_base() -> tf.keras.layers.Layer:
       inner_activation='relu')
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class XLNetMaskedLMTest(keras_parameterized.TestCase):
+class XLNetMaskedLMTest(tf.test.TestCase):
 
   def test_xlnet_masked_lm_head(self):
     hidden_size = 10
@@ -62,8 +58,7 @@ class XLNetMaskedLMTest(keras_parameterized.TestCase):
     self.assertAllClose(mlm_output.shape, (batch_size, hidden_size))
 
 
-@keras_parameterized.run_all_keras_modes
-class XLNetPretrainerTest(keras_parameterized.TestCase):
+class XLNetPretrainerTest(tf.test.TestCase):
 
   def test_xlnet_trainer(self):
     """Validates that the Keras object can be created."""
@@ -144,8 +139,7 @@ class XLNetPretrainerTest(keras_parameterized.TestCase):
                         new_xlnet_trainer_model.get_config())
 
 
-@keras_parameterized.run_all_keras_modes
-class XLNetClassifierTest(keras_parameterized.TestCase):
+class XLNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
 
   def test_xlnet_trainer(self):
     """Validate that the Keras object can be created."""
@@ -232,8 +226,7 @@ class XLNetClassifierTest(keras_parameterized.TestCase):
                         new_xlnet_trainer_model.get_config())
 
 
-@keras_parameterized.run_all_keras_modes
-class XLNetSpanLabelerTest(keras_parameterized.TestCase):
+class XLNetSpanLabelerTest(tf.test.TestCase):
 
   def test_xlnet_trainer(self):
     """Validate that the Keras object can be created."""

official/nlp/modeling/networks/albert_encoder_test.py

@@ -14,22 +14,14 @@
 
 """Tests for ALBERT transformer-based text encoder network."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import albert_encoder
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class AlbertEncoderTest(keras_parameterized.TestCase):
+class AlbertEncoderTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(AlbertEncoderTest, self).tearDown()

official/nlp/modeling/networks/bert_dense_encoder_test.py

@@ -19,14 +19,10 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import bert_encoder
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertEncoderV2Test(keras_parameterized.TestCase):
+class BertEncoderV2Test(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(BertEncoderV2Test, self).tearDown()

official/nlp/modeling/networks/bert_encoder_test.py

@@ -19,14 +19,10 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import bert_encoder
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class BertEncoderTest(keras_parameterized.TestCase):
+class BertEncoderTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(BertEncoderTest, self).tearDown()

official/nlp/modeling/networks/classification_test.py

@@ -14,22 +14,14 @@
 
 """Tests for classification network."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import classification
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class ClassificationTest(keras_parameterized.TestCase):
+class ClassificationTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(1, 10)
   def test_network_creation(self, num_classes):

official/nlp/modeling/networks/encoder_scaffold_test.py

@@ -18,7 +18,6 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.modeling import activations
 from official.nlp.modeling import layers
 from official.nlp.modeling.networks import encoder_scaffold
@@ -76,10 +75,7 @@ class TestLayer(tf.keras.layers.Layer):
   pass
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class EncoderScaffoldLayerClassTest(keras_parameterized.TestCase):
+class EncoderScaffoldLayerClassTest(tf.test.TestCase, parameterized.TestCase):
 
   def tearDown(self):
     super(EncoderScaffoldLayerClassTest, self).tearDown()
@@ -384,8 +380,7 @@ class Embeddings(tf.keras.Model):
     return word_embeddings, self.attention_mask([word_embeddings, mask])
 
 
-@keras_parameterized.run_all_keras_modes
-class EncoderScaffoldEmbeddingNetworkTest(keras_parameterized.TestCase):
+class EncoderScaffoldEmbeddingNetworkTest(tf.test.TestCase):
 
   def test_network_invocation(self):
     hidden_size = 32
@@ -525,8 +520,8 @@ class EncoderScaffoldEmbeddingNetworkTest(keras_parameterized.TestCase):
       new_network.get_embedding_table()
 
 
-@keras_parameterized.run_all_keras_modes
-class EncoderScaffoldHiddenInstanceTest(keras_parameterized.TestCase):
+class EncoderScaffoldHiddenInstanceTest(
+    tf.test.TestCase, parameterized.TestCase):
 
   def test_network_invocation(self):
     hidden_size = 32

official/nlp/modeling/networks/span_labeling_test.py

@@ -16,14 +16,10 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import span_labeling
 
 
-# This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
-# guarantees forward compatibility of this code for the V2 switchover.
-@keras_parameterized.run_all_keras_modes
-class SpanLabelingTest(keras_parameterized.TestCase):
+class SpanLabelingTest(tf.test.TestCase):
 
   def test_network_creation(self):
     """Validate that the Keras object can be created."""
@@ -165,8 +161,7 @@ class SpanLabelingTest(keras_parameterized.TestCase):
       _ = span_labeling.SpanLabeling(input_width=10, output='bad')
 
 
-@keras_parameterized.run_all_keras_modes
-class XLNetSpanLabelingTest(keras_parameterized.TestCase):
+class XLNetSpanLabelingTest(tf.test.TestCase):
 
   def test_basic_invocation_train(self):
     batch_size = 2
@@ -282,8 +277,8 @@ class XLNetSpanLabelingTest(keras_parameterized.TestCase):
 
     # Test `call` with training flag.
     # Note: this fails due to incompatibility with the functional API.
-    with self.assertRaisesRegexp(AssertionError,
-                                 'Could not compute output KerasTensor'):
+    with self.assertRaisesRegex(AssertionError,
+                                'Could not compute output KerasTensor'):
       model(inputs, training=True)
 
   def test_serialize_deserialize(self):

official/nlp/modeling/networks/xlnet_base_test.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 """Tests for Keras based XLNet model."""
+
+from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
 from tensorflow.python.distribute import combinations
-from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
 from official.nlp.modeling.networks import xlnet_base
 
 
-@keras_parameterized.run_all_keras_modes
-class RelativePositionEncodingTest(keras_parameterized.TestCase):
+class RelativePositionEncodingTest(tf.test.TestCase):
 
   def test_positional_embedding(self):
     """A low-dimensional example is tested.
@@ -47,7 +47,7 @@ class RelativePositionEncodingTest(keras_parameterized.TestCase):
     self.assertAllClose(encoding, target)
 
 
-class ComputePositionEncodingTest(keras_parameterized.TestCase):
+class ComputePositionEncodingTest(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       attention_type=["uni", "bi"],
@@ -111,7 +111,7 @@ class CausalAttentionMaskTests(tf.test.TestCase):
     self.assertAllClose(causal_attention_mask, expected_output)
 
 
-class MaskComputationTests(keras_parameterized.TestCase):
+class MaskComputationTests(tf.test.TestCase, parameterized.TestCase):
 
   @combinations.generate(combinations.combine(
       use_input_mask=[False, True],