Repository: keras-team/keras-core

Issue:
TensorFlow CuDnn GPU Test Failure - GRU & LSTM
When I run 

Open Terminal in Colab Pro with GPU:
```
git clone https://github.com/keras-team/keras-core.git
cd keras-core
# Updating TensorFlow in Colab to 2.13 version.
pip install -U tensorflow
# Run with TF backend for RNN (LSTM / Gru)
KERAS_BACKEND=tensorflow pytest keras_core/layers/rnn
```

Error:
```
E                                           tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:
E                                           
E                                           Detected unsupported operations when trying to compile graph __inference_one_step_on_data_31185[] on XLA_GPU_JIT: CudnnRNN (No registered 'CudnnRNN' OpKernel for XLA_GPU_JIT devices compatible with node {{node test_model_15_1/lstm_9/CudnnRNN}}){{node test_model_15_1/lstm_9/CudnnRNN}}
E                                           The op is created at: 
E                                           File "usr/local/bin/pytest", line 8, in <module>
E                                             sys.exit(console_main())
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py", line 190, in console_main
E                                             code = main()
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py", line 167, in main
E                                             ret: Union[ExitCode, int] = config.hook.pytest_cmdline_main(
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_hooks.py", line 433, in __call__
E                                             return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
E                                             return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
E                                             res = hook_impl.function(*args)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/main.py", line 317, in pytest_cmdline_main
E                                             return wrap_session(config, _main)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/main.py", line 270, in wrap_session
E                                             session.exitstatus = doit(config, session) or 0
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/main.py", line 324, in _main
E                                             config.hook.pytest_runtestloop(session=session)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_hooks.py", line 433, in __call__
E                                             return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
E                                             return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
E                                             res = hook_impl.function(*args)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/main.py", line 349, in pytest_runtestloop
E                                             item.config.hook.pytest_runtest_protocol(item=item, nextitem=nextitem)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_hooks.py", line 433, in __call__
E                                             return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
E                                             return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
E                                             res = hook_impl.function(*args)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 112, in pytest_runtest_protocol
E                                             runtestprotocol(item, nextitem=nextitem)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 131, in runtestprotocol
E                                             reports.append(call_and_report(item, "call", log))
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 220, in call_and_report
E                                             call = call_runtest_hook(item, when, **kwds)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 259, in call_runtest_hook
E                                             return CallInfo.from_call(
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 339, in from_call
E                                             result: Optional[TResult] = func()
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 260, in <lambda>
E                                             lambda: ihook(item=item, **kwds), when=when, reraise=reraise
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_hooks.py", line 433, in __call__
E                                             return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
E                                             return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
E                                           File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
E                                             res = hook_impl.function(*args)
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 167, in pytest_runtest_call
E                                             item.runtest()
E                                           File "usr/local/lib/python3.10/dist-packages/_pytest/unittest.py", line 330, in runtest
E                                             self._testcase(result=self)  # type: ignore[arg-type]
E                                           File "usr/lib/python3.10/unittest/case.py", line 650, in __call__
E                                             return self.run(*args, **kwds)
E                                           File "usr/lib/python3.10/unittest/case.py", line 591, in run
E                                             self._callTestMethod(testMethod)
E                                           File "usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
E                                             method()
E                                           File "content/keras-core/keras_core/layers/rnn/lstm_test.py", line 23, in test_basics
E                                             self.run_layer_test(
E                                           File "content/keras-core/keras_core/testing/test_case.py", line 334, in run_layer_test
E                                             run_training_step(layer, input_data, output_data)
E                                           File "content/keras-core/keras_core/testing/test_case.py", line 295, in run_training_step
E                                             model.fit(input_data, output_data, verbose=0)
E                                           File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
E                                             return fn(*args, **kwargs)
E                                           File "content/keras-core/keras_core/backend/tensorflow/trainer.py", line 323, in fit
E                                             logs = self.train_function(iterator)
E                                           File "content/keras-core/keras_core/backend/tensorflow/trainer.py", line 118, in one_step_on_iterator
E                                             outputs = self.distribute_strategy.run(
E                                           File "content/keras-core/keras_core/backend/tensorflow/trainer.py", line 106, in one_step_on_data
E                                             return self.train_step(data)
E                                           File "content/keras-core/keras_core/backend/tensorflow/trainer.py", line 61, in train_step
E                                             y_pred = self(x)
E                                           File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
E                                             return fn(*args, **kwargs)
E                                           File "content/keras-core/keras_core/layers/layer.py", line 738, in __call__
E                                             outputs = super().__call__(*args, **kwargs)
E                                           File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
E                                             return fn(*args, **kwargs)
E                                           File "content/keras-core/keras_core/ops/operation.py", line 46, in __call__
E                                             return self.call(*args, **kwargs)
E                                           File "content/keras-core/keras_core/testing/test_case.py", line 285, in call
E                                             return self.layer(x)
E                                           File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
E                                             return fn(*args, **kwargs)
E                                           File "content/keras-core/keras_core/layers/layer.py", line 738, in __call__
E                                             outputs = super().__call__(*args, **kwargs)
E                                           File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
E                                             return fn(*args, **kwargs)
E                                           File "content/keras-core/keras_core/ops/operation.py", line 46, in __call__
E                                             return self.call(*args, **kwargs)
E                                           File "content/keras-core/keras_core/layers/rnn/lstm.py", line 531, in call
E                                             return super().call(
E                                           File "content/keras-core/keras_core/layers/rnn/rnn.py", line 390, in call
E                                             last_output, outputs, states = self.inner_loop(
E                                           File "content/keras-core/keras_core/layers/rnn/lstm.py", line 505, in inner_loop
E                                             out = backend.lstm(
E                                           File "content/keras-core/keras_core/backend/tensorflow/rnn.py", line 814, in lstm
E                                             return _cudnn_lstm(
E                                           File "content/keras-core/keras_core/backend/tensorflow/rnn.py", line 925, in _cudnn_lstm
E                                             outputs, h, c, _ = tf.raw_ops.CudnnRNN(
E                                                [[StatefulPartitionedCall]] [Op:__inference_one_step_on_iterator_31204]

/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py:53: InvalidArgumentError
-------------------------------------------- Captured stderr call ---------------------------------------------
2023-07-26 04:44:21.180475: W tensorflow/core/framework/op_kernel.cc:1828] OP_REQUIRES failed at xla_ops.cc:503 : INVALID_ARGUMENT: Detected unsupported operations when trying to compile graph __inference_one_step_on_data_31185[] on XLA_GPU_JIT: CudnnRNN (No registered 'CudnnRNN' OpKernel for XLA_GPU_JIT devices compatible with node {{node test_model_15_1/lstm_9/CudnnRNN}}){{node test_model_15_1/lstm_9/CudnnRNN}}
The op is created at: 
File "usr/local/bin/pytest", line 8, in <module>
  sys.exit(console_main())
File "usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py", line 190, in console_main
  code = main()
File "usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py", line 167, in main
  ret: Union[ExitCode, int] = config.hook.pytest_cmdline_main(
File "usr/local/lib/python3.10/dist-packages/pluggy/_hooks.py", line 433, in __call__
  return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
  return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
  res = hook_impl.function(*args)

  return self._hookexec(self.name, self._hookimpls, kwargs, firstresult)
File "usr/local/lib/python3.10/dist-packages/pluggy/_manager.py", line 112, in _hookexec
  return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
File "usr/local/lib/python3.10/dist-packages/pluggy/_callers.py", line 80, in _multicall
  res = hook_impl.function(*args)
File "usr/local/lib/python3.10/dist-packages/_pytest/runner.py", line 167, in pytest_runtest_call
  item.runtest()
File "usr/local/lib/python3.10/dist-packages/_pytest/unittest.py", line 330, in runtest
  self._testcase(result=self)  # type: ignore[arg-type]
File "usr/lib/python3.10/unittest/case.py", line 650, in __call__
  return self.run(*args, **kwds)
File "usr/lib/python3.10/unittest/case.py", line 591, in run
  self._callTestMethod(testMethod)
File "usr/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
  method()
File "content/keras-core/keras_core/layers/rnn/lstm_test.py", line 23, in test_basics
  self.run_layer_test(
File "content/keras-core/keras_core/testing/test_case.py", line 334, in run_layer_test
  run_training_step(layer, input_data, output_data)
File "content/keras-core/keras_core/testing/test_case.py", line 295, in run_training_step
  model.fit(input_data, output_data, verbose=0)
File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
  return fn(*args, **kwargs)
File "content/keras-core/keras_core/ops/operation.py", line 46, in __call__
  return self.call(*args, **kwargs)
File "content/keras-core/keras_core/testing/test_case.py", line 285, in call
  return self.layer(x)
File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
  return fn(*args, **kwargs)
File "content/keras-core/keras_core/layers/layer.py", line 738, in __call__
  outputs = super().__call__(*args, **kwargs)
File "content/keras-core/keras_core/utils/traceback_utils.py", line 114, in error_handler
  return fn(*args, **kwargs)
File "content/keras-core/keras_core/ops/operation.py", line 46, in __call__
  return self.call(*args, **kwargs)
File "content/keras-core/keras_core/layers/rnn/lstm.py", line 531, in call
  return super().call(
File "content/keras-core/keras_core/layers/rnn/rnn.py", line 390, in call
  last_output, outputs, states = self.inner_loop(
File "content/keras-core/keras_core/layers/rnn/lstm.py", line 505, in inner_loop
  out = backend.lstm(
File "content/keras-core/keras_core/backend/tensorflow/rnn.py", line 814, in lstm
  return _cudnn_lstm(
File "content/keras-core/keras_core/backend/tensorflow/rnn.py", line 925, in _cudnn_lstm
  outputs, h, c, _ = tf.raw_ops.CudnnRNN(
=========================================== short test summary info ===========================================
FAILED keras_core/layers/rnn/gru_test.py::GRUTest::test_basics - tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:
FAILED keras_core/layers/rnn/lstm_test.py::LSTMTest::test_basics - tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:
```


Discussion:
> Detected unsupported operations when trying to compile graph __inference_one_step_on_data_31185[] on XLA_GPU_JIT: CudnnRNN (No registered 'CudnnRNN' OpKernel for XLA_GPU_JIT devices compatible with node {{node test_model_15_1/lstm_9/CudnnRNN}}){{node test_model_15_1/lstm_9/CudnnRNN}}

This comes from trying to use cuDNN ops with `jit_compile=True`. I believe I have fixed the default behavior already (by disabling `jit_compile` when using cuDNN RNN layers). I guess this happens because some tests pass `jit_compile` explicitly? We should fix that.
I found that there's actually a bug with the disabling of jit_compile, and it's subtle.
Bug should be fixed -- please check.

Resolving patch:
diff --git a/keras_core/trainers/trainer.py b/keras_core/trainers/trainer.py
--- a/keras_core/trainers/trainer.py
+++ b/keras_core/trainers/trainer.py
@@ -100,7 +100,14 @@ def jit_compile(self):
 
     @jit_compile.setter
     def jit_compile(self, value):
-        self._jit_compile = value
+        if value and not model_supports_jit(self):
+            warnings.warn(
+                "Model doesn't support `jit_compile=True`. "
+                "Proceeding with `jit_compile=False`."
+            )
+            self._jit_compile = False
+        else:
+            self._jit_compile = value
 
     @property
     def run_eagerly(self):


Test patch: