Code changes to get ready for an incoming Keras optimizer migration.

PiperOrigin-RevId: 472844636

Author: chenmoneygithub

tensorflow_probability/python/bijectors/masked_autoregressive_test.py

@@ -797,7 +797,7 @@ def test_doc_string_2(self):
         x_, bijector_kwargs={"conditional_input": c_})
     model = tfk.Model([x_, c_], log_prob_)
 
-    if tf.__internal__.tf2.enabled():
+    if tf.__internal__.tf2.enabled() and tf.executing_eagerly():
       optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)
     else:
       optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=0.1)

tensorflow_probability/python/distributions/lambertw_f_test.py

@@ -193,7 +193,7 @@ def dist_lambda(t):
         tf.keras.layers.Dense(1 + 1 + 1),
         dist_layer])
     negloglik = lambda y, p_y: -p_y.log_prob(y)
-    if tf.__internal__.tf2.enabled():
+    if tf.__internal__.tf2.enabled() and tf.executing_eagerly():
       optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
     else:
       optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=0.01)

tensorflow_probability/python/layers/dense_variational_v2_test.py

@@ -84,7 +84,7 @@ def test_end_to_end(self):
         tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1))
     ])
 
-    if tf.__internal__.tf2.enabled():
+    if tf.__internal__.tf2.enabled() and tf.executing_eagerly():
       optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
     else:
       optimizer = tf.keras.optimizers.legacy.Adam(learning_rate=0.05)

tensorflow_probability/python/layers/distribution_layer_test.py

@@ -55,7 +55,7 @@ def _unwrap_tensor_coercible(dist):
 
 
 def _get_adam_optimizer(learning_rate):
-  if tf.__internal__.tf2.enabled():
+  if tf.__internal__.tf2.enabled() and tf.executing_eagerly():
     return tf.keras.optimizers.Adam(learning_rate=learning_rate)
   return tf.keras.optimizers.legacy.Adam(learning_rate=learning_rate)
 

tensorflow_probability/python/layers/weight_norm_test.py

@@ -224,7 +224,7 @@ def testGradientValues(self, model_type):
 
   @parameterized.parameters(['sequential', 'sequential_no_input', 'functional'])
   def testTrainableVariableInitializationInModelFit(self, model_type):
-    if tf.__internal__.tf2.enabled():
+    if tf.__internal__.tf2.enabled() and tf.executing_eagerly():
       sgd = tf.keras.optimizers.SGD(lr=0.)
     else:
       sgd = tf.keras.optimizers.legacy.SGD(lr=0.)