Failed to determine best cudnn convolution algorithm when using mixed_float16 policy
@sachinprasadhs I was able to replicate the issue reported [here](https://colab.research.google.com/gist/sushreebarsa/6fdcee7ffcbc7773d8dcf8ff7b65524c/failed-to-determine-best-cudnn-convolution-algorithm.ipynb#scrollTo=KuFHU9kYs24S). Thank you! It seems to be a problem with XLA. When `jit_compile` is set to `False`, the mixed_float16 behaves as expected. @shkarupa-alex Could you please confirm if this solution works for you? After confirmation, we can investigate the underlying issue in XLA. @kaixih i confirm that issue raises only when precision policy is mixed_float16 and jit_compile=True.
In all other cases it works well.

But i need both mixed_float16 and jit_compile=True.
In my models i see that with jit_compile=True batch size could be twice larger then with autoclustering Certainly. We are investigating whether the problem lies in the jit compilation or if it is related to an issue with the Nvidia library. Rest assured, we are actively working on resolving this issue and will keep you informed of any progress. Thank you for your patience. There is another error message in TF 2.15 + cuda 12.2:
```
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPU will likely run quickly with dtype policy mixed_float16 as it has compute capability of at least 7.0. Your GPU: NVIDIA GeForce RTX 4090, compute capability 8.9
loc(loc(""ddoott..5587"")): : error: error: ''arith.mulfarith.mulf' op ' op requires the same type for all operands and resultsrequires the same type for all operands and results

F0000 00:00:1700554808.926726  189829 ir_emitter_triton.cc:1672] Check failed: mlir::succeeded(mlir::verify(*triton_module)) 
F0000 00:00:1700554808.926724  189834 ir_emitter_triton.cc:1672] Check failed: mlir::succeeded(mlir::verify(*triton_module)) 
```
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

'v2.13.0-rc2-7-g1cb1a030a62', '2.13.0'

### Custom code

Yes

### OS platform and distribution

Google Colab

### Mobile device

No

### Python version

Google Colab default

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

Google Colab default

### GPU model and memory

Google Colab default

### Current behavior?

When using depthwise conv inside custom attention (https://arxiv.org/pdf/2304.04237.pdf) layer with **mixed_float16** policy got "Failed to determine best cudnn convolution algorithm".

There is no such error with full precision.

### Standalone code to reproduce the issue

```shell
https://colab.research.google.com/drive/1vsdCr_Mus9N6xax1vKjDvN093dh5F6rO?usp=sharing
```


### Relevant log output

```shell
---------------------------------------------------------------------------
UnknownError                              Traceback (most recent call last)
<ipython-input-3-32ea805343ae> in <cell line: 23>()
     21 # Train
     22 model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', run_eagerly=False, jit_compile=True)
---> 23 model.fit(ds_train, epochs=1, steps_per_epoch=20)

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     51   try:
     52     ctx.ensure_initialized()
---> 53     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     54                                         inputs, attrs, num_outputs)
     55   except core._NotOkStatusException as e:

UnknownError: Graph execution error:

Detected at node 'StatefulPartitionedCall' defined at (most recent call last):
    File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.10/dist-packages/colab_kernel_launcher.py", line 37, in <module>
      ColabKernelApp.launch_instance()
    File "/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py", line 619, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py", line 195, in start
      self.asyncio_loop.run_forever()
    File "/usr/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
      self._run_once()
    File "/usr/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
      handle._run()
    File "/usr/lib/python3.10/asyncio/events.py", line 80, in _run
      self._context.run(self._callback, *self._args)
    File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 685, in <lambda>
      lambda f: self._run_callback(functools.partial(callback, future))
    File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 738, in _run_callback
      ret = callback()
    File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 825, in inner
      self.ctx_run(self.run)
    File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 786, in run
      yielded = self.gen.send(value)
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 361, in process_one
      yield gen.maybe_future(dispatch(*args))
    File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell
      yield gen.maybe_future(handler(stream, idents, msg))
    File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request
      self.do_execute(
    File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
      yielded = ctx_run(next, result)
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute
      res = shell.run_cell(code, store_history=store_history, silent=silent)
    File "/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell
      return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell
      result = self._run_cell(
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell
      return runner(coro)
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner
      coro.send(None)
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async
      has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes
      if (await self.run_code(code, result,  async_=asy)):
    File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code
      exec(code_obj, self.user_global_ns, self.user_ns)
    File "<ipython-input-3-32ea805343ae>", line 23, in <cell line: 23>
      model.fit(ds_train, epochs=1, steps_per_epoch=20)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1742, in fit
      tmp_logs = self.train_function(iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1338, in train_function
      return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1322, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
Node: 'StatefulPartitionedCall'
Failed to determine best cudnn convolution algorithm for:
%cudnn-conv.19 = (f16[8,128,128,64]{3,2,1,0}, u8[0]{0}) custom-call(f16[8,128,128,1600]{3,2,1,0} %bitcast.959, f16[64,5,5,32]{3,2,1,0} %pad.15), window={size=5x5 pad=2_2x2_2}, dim_labels=b01f_o01i->b01f, feature_group_count=64, custom_call_target="__cudnn$convForward", metadata={op_type="DepthwiseConv2dNativeBackpropInput" op_name="gradient_tape/model_1/slide_attention_1/depthwise/DepthwiseConv2dNativeBackpropInput" source_file="/usr/local/lib/python3.10/dist-packages/keras/src/optimizers/optimizer.py" source_line=276}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: UNKNOWN: CUDNN_STATUS_BAD_PARAM
in tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc(3999): 'op' CUDNN_BACKEND_OPERATION: cudnnFinalize Failed

To ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.
	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_train_function_4989]
```
