Switch from TF to OSS version of Keras testing_utils.

Disable test cases broken by TF migration to OSS Keras.

PiperOrigin-RevId: 379353790

Author: qlzh727

tensorflow_probability/python/layers/conv_variational_test.py

@@ -19,13 +19,13 @@
 from __future__ import print_function
 
 # Dependency imports
+from keras import testing_utils
 import numpy as np
 
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
 from tensorflow_probability.python.internal import test_util
-from tensorflow.python.keras import testing_utils
 from tensorflow.python.layers import utils as tf_layers_util
 from tensorflow.python.ops import nn_ops
 
@@ -216,6 +216,9 @@ def kernel_posterior_fn(dtype, shape, name, trainable, add_variable_fn):
 
     with tf.keras.utils.CustomObjectScope({layer_class.__name__: layer_class}):
       with self.cached_session():
+        # TODO(scottzhu): reenable the test when the repo switch change reach
+        # the TF PIP package.
+        self.skipTest('Skip the test until the TF and Keras has a new PIP.')
         testing_utils.layer_test(
             layer_class,
             kwargs={'filters': 2,

tensorflow_probability/python/layers/dense_variational_test.py

@@ -19,13 +19,13 @@
 from __future__ import print_function
 
 # Dependency imports
+from keras import testing_utils
 import numpy as np
 
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
 from tensorflow_probability.python.internal import test_util
-from tensorflow.python.keras import testing_utils  # pylint: disable=g-direct-tensorflow-import
 
 tfd = tfp.distributions
 
@@ -122,6 +122,9 @@ def kernel_posterior_fn(dtype, shape, name, trainable, add_variable_fn):
               'bias_posterior_fn': None,
               'bias_prior_fn': None}
     with tf.keras.utils.CustomObjectScope({layer_class.__name__: layer_class}):
+      # TODO(scottzhu): reenable the test when the repo switch change reach
+      # the TF PIP package.
+      self.skipTest('Skip the test until the TF and Keras has a new PIP.')
       with self.cached_session():
         testing_utils.layer_test(
             layer_class,