Merge branch 'keras-team:master' into fix_tpu_tests

Author: divyashreepathihalli

.kokoro/github/ubuntu/gpu/jax/continuous.cfg

@@ -12,5 +12,5 @@ env_vars: {
    value: "jax"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 120 mins from default 180 mins
+timeout_mins: 120

.kokoro/github/ubuntu/gpu/jax/presubmit.cfg

@@ -12,5 +12,5 @@ env_vars: {
    value: "jax"
 }
 
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
+# Set timeout to 120 mins from default 180 mins
+timeout_mins: 120

keras/src/export/litert_test.py

@@ -36,8 +36,11 @@
 
 # Model types to test (LSTM only if AI Edge LiteRT is available)
 model_types = ["sequential", "functional"]
-if AI_EDGE_LITERT_AVAILABLE:
-    model_types.append("lstm")
+# TODO(#21914): `"lstm"` does not work with ai-edge-litert==1.3.0.
+# Unfortunately, for TF 2.20.0, this is the only version which works. Uncomment
+# this part when we upgrade TF and ai-edge-litert.
+# if AI_EDGE_LITERT_AVAILABLE:
+#     model_types.append("lstm")
 
 
 class CustomModel(models.Model):

requirements-tensorflow-cuda.txt

@@ -1,7 +1,9 @@
 # Tensorflow with cuda support.
 tensorflow[and-cuda]~=2.20.0
 tf2onnx
-ai-edge-litert
+# This is the only version which works with TF 2.20.0.
+# TODO(#21914): Update this version when TF is updated.
+ai-edge-litert==1.3.0
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu

requirements.txt

@@ -4,7 +4,9 @@
 tensorflow-cpu~=2.20.0;sys_platform != 'darwin'
 tensorflow~=2.20.0;sys_platform == 'darwin'
 tf2onnx
-ai-edge-litert
+# This is the only version which works with TF 2.20.0.
+# TODO(#21914): Update this version when TF is updated.
+ai-edge-litert==1.3.0
 
 # Torch.
 --extra-index-url https://download.pytorch.org/whl/cpu