Fix tf2.4 (#9120)

* Fix tests for TF 2.4

* Remove <2.4 limitation

* Add version condition

* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>

setup.py

@@ -127,8 +127,8 @@ _deps = [
     "sphinx-rtd-theme==0.4.3",  # sphinx-rtd-theme==0.5.0 introduced big changes in the style.
     "sphinx==3.2.1",
     "starlette",
-    "tensorflow-cpu>=2.0,<2.4",
-    "tensorflow>=2.0,<2.4",
+    "tensorflow-cpu>=2.0",
+    "tensorflow>=2.0",
     "timeout-decorator",
     "tokenizers==0.9.4",
     "torch>=1.0",

tests/test_modeling_tf_common.py

@@ -434,14 +434,14 @@ class TFModelTesterMixin:
                 num_labels = 2
 
             X = tf.data.Dataset.from_tensor_slices(
-                (inputs_dict, np.random.randint(0, num_labels, (self.model_tester.batch_size, 1)))
+                (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)))
             ).batch(1)
 
             hidden_states = main_layer(symbolic_inputs)[0]
             outputs = tf.keras.layers.Dense(num_labels, activation="softmax", name="outputs")(hidden_states)
             model = tf.keras.models.Model(inputs=symbolic_inputs, outputs=[outputs])
 
-            model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"])
+            model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["binary_accuracy"])
             model.fit(X, epochs=1)
 
             with tempfile.TemporaryDirectory() as tmpdirname:

tests/test_optimization_tf.py

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