Use stable functions (#9369)

tests/test_modeling_tf_common.py

@@ -51,10 +51,10 @@ if is_tf_available():
         for gpu in gpus:
             # Restrict TensorFlow to only allocate x GB of memory on the GPUs
             try:
-                tf.config.experimental.set_virtual_device_configuration(
-                    gpu, [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
+                tf.config.set_logical_device_configuration(
+                    gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
                 )
-                logical_gpus = tf.config.experimental.list_logical_devices("GPU")
+                logical_gpus = tf.config.list_logical_devices("GPU")
                 print("Logical GPUs", logical_gpus)
             except RuntimeError as e:
                 # Virtual devices must be set before GPUs have been initialized

tests/test_optimization_tf.py

@@ -14,8 +14,6 @@
 
 import unittest
 
-from packaging import version
-
 from transformers import is_tf_available
 from transformers.testing_utils import require_tf
 
@@ -78,18 +76,12 @@ class OptimizationFTest(unittest.TestCase):
                 local_variables = strategy.experimental_local_results(gradient_placeholder)
                 local_variables[0].assign(grad1)
                 local_variables[1].assign(grad2)
-                if version.parse(tf.version.VERSION) >= version.parse("2.2"):
-                    strategy.run(accumulate_on_replica, args=(gradient_placeholder,))
-                else:
-                    strategy.experimental_run_v2(accumulate_on_replica, args=(gradient_placeholder,))
+                strategy.run(accumulate_on_replica, args=(gradient_placeholder,))
 
         @tf.function
         def apply_grad():
             with strategy.scope():
-                if version.parse(tf.version.VERSION) >= version.parse("2.2"):
-                    strategy.run(apply_on_replica)
-                else:
-                    strategy.experimental_run_v2(apply_on_replica)
+                strategy.run(apply_on_replica)
 
         def _check_local_values(grad1, grad2):
             values = strategy.experimental_local_results(accumulator._gradients[0])