Add tf_keras imports to prepare for Keras 3 (#28588)

* Port core files + ESM (because ESM code is odd)

* Search-replace in modelling code

* Fix up transfo_xl as well

* Fix other core files + tests (still need to add correct import to tests)

* Fix cookiecutter

* make fixup, fix imports in some more core files

* Auto-add imports to tests

* Cleanup, add imports to sagemaker tests

* Use correct exception for importing tf_keras

* Fixes in modeling_tf_utils

* make fixup

* Correct version parsing code

* Ensure the pipeline tests correctly revert to float32 after each test

* Ensure the pipeline tests correctly revert to float32 after each test

* More tf.keras -> keras

* Add dtype cast

* Better imports of tf_keras

* Add a cast for tf.assign, just in case

* Fix callback imports

tests/utils/test_modeling_tf_core.py

@@ -46,6 +46,7 @@ if is_tf_available():
         TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
         TFSharedEmbeddings,
     )
+    from transformers.modeling_tf_utils import keras
 
     if _tf_gpu_memory_limit is not None:
         gpus = tf.config.list_physical_devices("GPU")
@@ -169,7 +170,7 @@ class TFCoreModelTesterMixin:
                 self.assertGreater(len(inputs_minus_labels), 0)
 
                 # Make sure it works with XLA!
-                model.compile(optimizer=tf.keras.optimizers.SGD(0.0), jit_compile=True)
+                model.compile(optimizer=keras.optimizers.SGD(0.0), jit_compile=True)
                 # Make sure the model fits without crashing regardless of where we pass the labels
                 history = model.fit(
                     prepared_for_class,
@@ -186,7 +187,7 @@ class TFCoreModelTesterMixin:
 
                 # Now test it with separate labels, to make sure that path works in XLA too.
                 model = model_class(config)
-                model.compile(optimizer=tf.keras.optimizers.SGD(0.0), jit_compile=True)
+                model.compile(optimizer=keras.optimizers.SGD(0.0), jit_compile=True)
                 history = model.fit(
                     inputs_minus_labels,
                     labels,
@@ -234,7 +235,7 @@ class TFCoreModelTesterMixin:
             with tempfile.TemporaryDirectory() as tmpdirname:
                 model.save_pretrained(tmpdirname, saved_model=True)
                 saved_model_dir = os.path.join(tmpdirname, "saved_model", "1")
-                model = tf.keras.models.load_model(saved_model_dir)
+                model = keras.models.load_model(saved_model_dir)
                 outputs = model(class_inputs_dict)
 
                 if self.is_encoder_decoder:
@@ -264,7 +265,7 @@ class TFCoreModelTesterMixin:
 
     @slow
     def test_mixed_precision(self):
-        tf.keras.mixed_precision.set_global_policy("mixed_float16")
+        keras.mixed_precision.set_global_policy("mixed_float16")
 
         # try/finally block to ensure subsequent tests run in float32
         try:
@@ -276,7 +277,7 @@ class TFCoreModelTesterMixin:
 
                 self.assertIsNotNone(outputs)
         finally:
-            tf.keras.mixed_precision.set_global_policy("float32")
+            keras.mixed_precision.set_global_policy("float32")
 
     @slow
     def test_train_pipeline_custom_model(self):
@@ -296,7 +297,7 @@ class TFCoreModelTesterMixin:
             if module_member_name.endswith("MainLayer")
             for module_member in (getattr(module, module_member_name),)
             if isinstance(module_member, type)
-            and tf.keras.layers.Layer in module_member.__bases__
+            and keras.layers.Layer in module_member.__bases__
             and getattr(module_member, "_keras_serializable", False)
         }
 
@@ -311,7 +312,7 @@ class TFCoreModelTesterMixin:
                 main_layer = main_layer_class(config)
 
             symbolic_inputs = {
-                name: tf.keras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
+                name: keras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
             }
 
             if hasattr(self.model_tester, "num_labels"):
@@ -324,8 +325,8 @@ class TFCoreModelTesterMixin:
             ).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])
+            outputs = keras.layers.Dense(num_labels, activation="softmax", name="outputs")(hidden_states)
+            model = keras.models.Model(inputs=symbolic_inputs, outputs=[outputs])
 
             model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["binary_accuracy"])
             model.fit(X, epochs=1)
@@ -334,7 +335,7 @@ class TFCoreModelTesterMixin:
                 filepath = os.path.join(tmpdirname, "keras_model.h5")
                 model.save(filepath)
                 if "T5" in main_layer_class.__name__:
-                    model = tf.keras.models.load_model(
+                    model = keras.models.load_model(
                         filepath,
                         custom_objects={
                             main_layer_class.__name__: main_layer_class,
@@ -342,10 +343,10 @@ class TFCoreModelTesterMixin:
                         },
                     )
                 else:
-                    model = tf.keras.models.load_model(
+                    model = keras.models.load_model(
                         filepath, custom_objects={main_layer_class.__name__: main_layer_class}
                     )
-                assert isinstance(model, tf.keras.Model)
+                assert isinstance(model, keras.Model)
                 model(inputs_dict)
 
     @slow