Repository: keras-team/keras

Issue:
Execution time in model.fit
i got issue when i try to train my model using keras 3.0. the steps per epoch take really long time when i run it on keras 3.0 compared to when i'm using keras_core. Is there anyway to fix it ?
Here is the screenshot on keras 3.0 
![image](https://github.com/keras-team/keras/assets/72650027/4338f3f9-9a5c-4e57-a40d-7832a1b2c521)

and this one on keras-core
![image](https://github.com/keras-team/keras/assets/72650027/7e893ec5-4157-4114-9d83-01a90ef03e5e)

layers used in this model are : 
```py
def create_model():
    model = keras.Sequential([
        Input((778, 1)),
        Conv1D(96, 1, activation = 'relu'),
        Conv1D(352, 1,  activation = 'relu'),
        Conv1D(224, 1,  activation = 'relu'),
        Conv1D(160, 1,  activation = 'relu'),
        Conv1D(320, 1,  activation = 'relu'),
        Bidirectional(LSTM(192, return_sequences = True, activation = 'tanh')),
        Bidirectional(LSTM(384, return_sequences = True, activation = 'tanh')),
        Bidirectional(LSTM(160, return_sequences = True, activation = 'tanh')),
        Dropout(0.2),


        Dense(4, activation = 'softmax')
    ])
    model.compile(loss = 'sparse_categorical_crossentropy', optimizer = 'adam', metrics = ['accuracy'])
    return model
```
i run the model on google collab. is there any way to fix it ? Thanks


Discussion:
This sort of issue is usually because you're not using the GPU.

1. What backend is this?
2. Can you check if the backend can detect the GPU?
i'm using tensorflow backend.
this is the way i detect the gpu :
![image](https://github.com/keras-team/keras/assets/72650027/161fd7ea-a997-4564-a948-540a6403765f)
When i train the model, it's also utilize the GPU memory.
I can reproduce the issue. The issue is specifically with the LSTM layers. Most likely the problem is that cuDNN is not being used/detected by Keras 3 in this case. Will look into it.

Resolving patch:
diff --git a/keras/backend/tensorflow/rnn.py b/keras/backend/tensorflow/rnn.py
--- a/keras/backend/tensorflow/rnn.py
+++ b/keras/backend/tensorflow/rnn.py
@@ -441,7 +441,6 @@ def _step(time, output_ta_t, *states):
     return last_output, outputs, new_states
 
 
-@tf.function
 def gru(
     inputs,
     initial_state,
@@ -457,7 +456,6 @@ def gru(
     time_major=False,
     reset_after=True,
 ):
-    inputs_supported = _do_rnn_inputs_support_cudnn(mask, time_major)
     cudnn_supported = cudnn_ok(
         activation,
         recurrent_activation,
@@ -465,7 +463,7 @@ def gru(
         use_bias=bias is not None,
         reset_after=reset_after,
     )
-    if not cudnn_supported or not inputs_supported:
+    if not cudnn_supported or mask is not None:
         raise NotImplementedError
 
     from keras.backend.tensorflow import Variable
@@ -535,21 +533,6 @@ def _do_lstm_arguments_support_cudnn(
     )
 
 
-def _do_rnn_inputs_support_cudnn(mask, time_major):
-    if tf.sysconfig.get_build_info()["is_rocm_build"]:
-        if mask is not None:
-            return tf.reduce_all(mask)
-        return True
-    if mask is None:
-        return True
-    if time_major:
-        mask = tf.transpose(mask)
-    return tf.logical_and(
-        _is_sequence_right_padded(mask),
-        tf.logical_not(_has_fully_masked_sequence(mask)),
-    )
-
-
 def _is_sequence_right_padded(mask):
     """Check the mask tensor and see if it right padded.
 
@@ -655,7 +638,6 @@ def _is_gpu_available():
     return bool(tf.config.list_logical_devices("GPU"))
 
 
-@tf.function(autograph=False)
 def _cudnn_gru(
     inputs,
     initial_state,
@@ -802,7 +784,6 @@ def cudnn_ok(
     return args_supported and _is_gpu_available()
 
 
-@tf.function
 def lstm(
     inputs,
     initial_state_h,
@@ -818,11 +799,10 @@ def lstm(
     unroll=False,
     time_major=False,
 ):
-    inputs_supported = _do_rnn_inputs_support_cudnn(mask, time_major)
     cudnn_supported = cudnn_ok(
         activation, recurrent_activation, unroll, use_bias=bias is not None
     )
-    if not cudnn_supported or not inputs_supported:
+    if not cudnn_supported or mask is not None:
         raise NotImplementedError
 
     from keras.backend.tensorflow import Variable
@@ -855,7 +835,6 @@ def lstm(
         raise NotImplementedError
 
 
-@tf.function(autograph=False)
 def _cudnn_lstm(
     inputs,
     initial_state_h,


Test patch: