Switch to non-experimental mixed precision API. (#2657)

Author: reedwm

tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py

@@ -200,9 +200,7 @@ def test_invalid_save_freq(tmp_path):
 def test_loss_scale_optimizer(tmp_path):
     test_model_filepath = str(tmp_path / "test_model.{epoch:02d}.h5")
     moving_avg = MovingAverage(tf.keras.optimizers.SGD(lr=2.0), average_decay=0.5)
-    optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
-        moving_avg, "dynamic"
-    )
+    optimizer = tf.keras.mixed_precision.LossScaleOptimizer(moving_avg)
     x, y, model = get_data_and_model(optimizer)
     save_freq = "epoch"
     avg_model_ckpt = AverageModelCheckpoint(

tensorflow_addons/layers/tests/noisy_dense_test.py

@@ -20,7 +20,6 @@
 
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.keras.mixed_precision.experimental import Policy
 
 from tensorflow_addons.utils import test_utils
 from tensorflow_addons.layers.noisy_dense import NoisyDense
@@ -47,7 +46,7 @@ def test_noisy_dense_dtype(dtype):
 @pytest.mark.usefixtures("maybe_run_functions_eagerly")
 def test_noisy_dense_with_policy():
     inputs = tf.convert_to_tensor(np.random.randint(low=0, high=7, size=(2, 2)))
-    layer = NoisyDense(5, dtype=Policy("mixed_float16"), name="noisy_dense_policy")
+    layer = NoisyDense(5, dtype="mixed_float16", name="noisy_dense_policy")
     outputs = layer(inputs)
     output_signature = layer.compute_output_signature(
         tf.TensorSpec(dtype="float16", shape=(2, 2))

tensorflow_addons/layers/tests/stochastic_depth_test.py

@@ -41,7 +41,7 @@ def stochastic_depth_test(seed, training):
 
 @pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_with_mixed_precision_policy():
-    policy = tf.keras.mixed_precision.experimental.global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
 
     shortcut = np.asarray([[0.2, 0.1, 0.4]])
     residual = np.asarray([[0.2, 0.4, 0.5]])

tensorflow_addons/seq2seq/tests/attention_wrapper_test.py

@@ -398,7 +398,7 @@ def _test_with_attention(
                 seed=1337
             )
 
-    policy = tf.keras.mixed_precision.experimental.global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
     sampler = sampler_py.TrainingSampler()
     my_decoder = basic_decoder.BasicDecoder(cell=cell, sampler=sampler)
     initial_state = cell.get_initial_state(
@@ -541,7 +541,7 @@ def set_random_state_for_tf_and_np():
 @pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_bahdanau_not_normalized():
     set_random_state_for_tf_and_np()
-    policy = tf.keras.mixed_precision.experimental.global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
     create_attention_mechanism = wrapper.BahdanauAttention
     create_attention_kwargs = {"kernel_initializer": "ones"}
     expected_final_output = basic_decoder.BasicDecoderOutput(
@@ -617,7 +617,7 @@ def test_bahdanau_normalized():
 @pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_luong_not_normalized():
     set_random_state_for_tf_and_np()
-    policy = tf.keras.mixed_precision.experimental.global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
     create_attention_mechanism = wrapper.LuongAttention
 
     expected_final_output = basic_decoder.BasicDecoderOutput(

tensorflow_addons/seq2seq/tests/decoder_test.py

@@ -181,7 +181,7 @@ def test_dynamic_decode_rnn_with_training_helper_matches_dynamic_rnn(
 
 @pytest.mark.usefixtures("run_with_mixed_precision_policy")
 def test_dynamic_decode_rnn_with_scheduled_embedding_training_sampler():
-    policy = tf.keras.mixed_precision.experimental.global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
     sequence_length = [3, 4, 3, 1]
     batch_size = 4
     input_depth = 7

tensorflow_addons/utils/test_utils.py

@@ -114,9 +114,9 @@ def _restore_py_ops_value():
 
 @pytest.fixture(scope="function", params=["float32", "mixed_float16"])
 def run_with_mixed_precision_policy(request):
-    tf.keras.mixed_precision.experimental.set_policy(request.param)
+    tf.keras.mixed_precision.set_global_policy(request.param)
     yield
-    tf.keras.mixed_precision.experimental.set_policy("float32")
+    tf.keras.mixed_precision.set_global_policy("float32")
 
 
 @pytest.fixture(scope="function", params=["channels_first", "channels_last"])